From 6c83e40883e903063d3313df949a1452a305919e Mon Sep 17 00:00:00 2001 From: cxjyxx_me <498731903@qq.com> Date: Thu, 14 Jan 2021 13:09:20 +0800 Subject: [PATCH] fix --- python/jittor/misc.py | 2 +- python/jittor/test/test_memory_profiler.py | 18 +++++++++--------- src/executor.cc | 9 +++++++-- src/executor.h | 7 +++++-- src/mem/allocator/temp_allocator.cc | 2 +- src/mem/allocator/temp_allocator.h | 2 +- src/memory_profiler.cc | 17 ++++++++++++++++- src/memory_profiler.h | 7 +++++-- src/opt/gopt/setitem_gopt.cc | 1 - src/pybind/py_var_tracer.cc | 9 ++++++--- 10 files changed, 51 insertions(+), 23 deletions(-) diff --git a/python/jittor/misc.py b/python/jittor/misc.py index 5f725476..2345a864 100644 --- a/python/jittor/misc.py +++ b/python/jittor/misc.py @@ -3,7 +3,7 @@ # Maintainers: # Dun Liang . # Wenyang Zhou <576825820@qq.com> -# +# Guoye Yang <498731903@qq.com> # # This file is subject to the terms and conditions defined in # file 'LICENSE.txt', which is part of this source code package. diff --git a/python/jittor/test/test_memory_profiler.py b/python/jittor/test/test_memory_profiler.py index 02085a93..3eb59b5d 100644 --- a/python/jittor/test/test_memory_profiler.py +++ b/python/jittor/test/test_memory_profiler.py @@ -1,9 +1,9 @@ # *************************************************************** -# Copyright (c) 2020 Jittor. Authors: -# Guowei Yang <471184555@qq.com> -# Meng-Hao Guo +# Copyright (c) 2021 Jittor. All Rights Reserved. +# Maintainers: +# Guoye Yang <498731903@qq.com> # Dun Liang . -# All Rights Reserved. +# # This file is subject to the terms and conditions defined in # file 'LICENSE.txt', which is part of this source code package. # *************************************************************** @@ -53,7 +53,7 @@ class TestMemoryProfiler(unittest.TestCase): jt.seed(seed) @unittest.skipIf(not jt.has_cuda, "Cuda not found") - @jt.flag_scope(use_cuda=1, use_stat_allocator=1, trace_py_var=2) + @jt.flag_scope(use_cuda=1, use_stat_allocator=1, trace_py_var=3, profile_memory_enable=1) def test_resnet(self): self.setup_seed(1) loss_list=[] @@ -85,13 +85,13 @@ class TestMemoryProfiler(unittest.TestCase): _, out = jt.get_max_memory_treemap() out_ = out.split('\n') assert(out_[0] == 'root()') - assert(out_[3] == ' ├─mnist_net(MnistNet)') - assert(out_[7] == ' | └─model(ResNet)') + assert(out_[3] == ' ├─/home/cxjyxx_me/miniconda3/lib/python3.8/runpy.py:194(_run_module_as_main)') + assert(out_[7] == ' | └─/home/cxjyxx_me/miniconda3/lib/python3.8/runpy.py:87(_run_code)') _, out = jt.get_max_memory_treemap(build_by=1) out_ = out.split('\n') assert(out_[0] == 'root()') - assert(out_[4] == ' ├─mnist_net(MnistNet)') - assert(out_[8] == ' | └─model(ResNet)') + assert(out_[4] == ' ├─/home/cxjyxx_me/miniconda3/lib/python3.8/runpy.py:194(_run_module_as_main)') + assert(out_[8] == ' | └─/home/cxjyxx_me/miniconda3/lib/python3.8/runpy.py:87(_run_code)') if __name__ == "__main__": unittest.main() diff --git a/src/executor.cc b/src/executor.cc index b80905ae..3f1b87f4 100644 --- a/src/executor.cc +++ b/src/executor.cc @@ -1,6 +1,9 @@ // *************************************************************** // Copyright (c) 2021 Jittor. All Rights Reserved. -// Maintainers: Dun Liang . +// Maintainers: +// Dun Liang . +// Guoye Yang <498731903@qq.com> +// // This file is subject to the terms and conditions defined in // file 'LICENSE.txt', which is part of this source code package. // *************************************************************** @@ -29,6 +32,7 @@ namespace jittor { Executor exe; extern MemoryProfiler memory_profiler; +DECLARE_FLAG(int, profile_memory_enable); // from fetch_op.cc extern list fetcher_to_free; @@ -424,7 +428,8 @@ void Executor::run_sync(vector vars, bool device_sync) { for (auto* var : op->outputs()) { var->alloc(allocator); } - memory_profiler.check(); + if (PREDICT_BRANCH_NOT_TAKEN(profile_memory_enable)) + memory_profiler.check(); LOGvvv << "Run" << op << "inputs:" << op->inputs() << "outputs:" << op->outputs(); op->do_prepare(jkl); bool is_cuda = op->flags.get(NodeFlags::_cuda); diff --git a/src/executor.h b/src/executor.h index cc5d4eb2..5c924e98 100644 --- a/src/executor.h +++ b/src/executor.h @@ -1,6 +1,9 @@ // *************************************************************** -// Copyright (c) 2021 Jittor. All Rights Reserved. -// Maintainers: Dun Liang . +// Copyright (c) 2021 Jittor. All Rights Reserved. +// Maintainers: +// Dun Liang . +// Guoye Yang <498731903@qq.com> +// // This file is subject to the terms and conditions defined in // file 'LICENSE.txt', which is part of this source code package. // *************************************************************** diff --git a/src/mem/allocator/temp_allocator.cc b/src/mem/allocator/temp_allocator.cc index 88c1398c..b6b2c5a9 100644 --- a/src/mem/allocator/temp_allocator.cc +++ b/src/mem/allocator/temp_allocator.cc @@ -1,5 +1,5 @@ // *************************************************************** -// Copyright (c) 2020 Jittor. All Rights Reserved. +// Copyright (c) 2021 Jittor. All Rights Reserved. // Maintainers: // Guoye Yang <498731903@qq.com> // Dun Liang . diff --git a/src/mem/allocator/temp_allocator.h b/src/mem/allocator/temp_allocator.h index 0402e421..50a0969b 100644 --- a/src/mem/allocator/temp_allocator.h +++ b/src/mem/allocator/temp_allocator.h @@ -1,5 +1,5 @@ // *************************************************************** -// Copyright (c) 2020 Jittor. All Rights Reserved. +// Copyright (c) 2021 Jittor. All Rights Reserved. // Maintainers: // Guoye Yang <498731903@qq.com> // Dun Liang . diff --git a/src/memory_profiler.cc b/src/memory_profiler.cc index 27dbcd8a..572ff213 100644 --- a/src/memory_profiler.cc +++ b/src/memory_profiler.cc @@ -1,3 +1,12 @@ +// *************************************************************** +// Copyright (c) 2021 Jittor. All Rights Reserved. +// Maintainers: +// Guoye Yang <498731903@qq.com> +// Dun Liang . +// +// This file is subject to the terms and conditions defined in +// file 'LICENSE.txt', which is part of this source code package. +// *************************************************************** #include "memory_profiler.h" #include "graph.h" #include "var_holder.h" @@ -19,7 +28,7 @@ struct FloatOutput_ { string suffix; int p=4; }; -std::ostream& operator<<(std::ostream& os, const FloatOutput_& o) { +inline std::ostream& operator<<(std::ostream& os, const FloatOutput_& o) { int w = 8; os << std::setw(w-2-o.suffix.size()); os << std::setprecision(o.p); @@ -47,6 +56,7 @@ void MemoryProfiler::clear() { } std::pair MemoryProfiler::get_memory_info() { + ASSERT(profile_memory_enable == 1); size_t used = 0; size_t unused = 0; //TODO add mssfrl allocator @@ -58,6 +68,7 @@ std::pair MemoryProfiler::get_memory_info() { } void MemoryProfiler::check() { + ASSERT(profile_memory_enable == 1); std::pair mem_info = get_memory_info(); if (mem_info.first > max_used_memory_size) { max_used_memory_size = mem_info.first; @@ -99,6 +110,7 @@ bool MemoryProfiler::cmp(const std::pair>, size_ } void MemoryProfiler::display_max_memory_info() { + ASSERT(profile_memory_enable == 1); Log log("", 'i', 0); std::sort(max_live_vars.begin(), max_live_vars.end(), cmp); log << "\n=====display_max_memory_info=====\n"; @@ -117,10 +129,12 @@ void MemoryProfiler::display_max_memory_info() { } void display_max_memory_info() { + ASSERT(profile_memory_enable == 1); memory_profiler.display_max_memory_info(); } string MemoryProfiler::get_max_memory_info() { + ASSERT(profile_memory_enable == 1); std::stringstream out; string div1 = "[!@#div1!@#]"; string div2 = "[!@#div2!@#]"; @@ -142,6 +156,7 @@ string MemoryProfiler::get_max_memory_info() { } string get_max_memory_info() { + ASSERT(profile_memory_enable == 1); return memory_profiler.get_max_memory_info(); } diff --git a/src/memory_profiler.h b/src/memory_profiler.h index 7a0ce12b..be1daac2 100644 --- a/src/memory_profiler.h +++ b/src/memory_profiler.h @@ -1,6 +1,9 @@ // *************************************************************** -// Copyright (c) 2020 Jittor. All Rights Reserved. -// Authors: Dun Liang . +// Copyright (c) 2021 Jittor. All Rights Reserved. +// Maintainers: +// Guoye Yang <498731903@qq.com> +// Dun Liang . +// // This file is subject to the terms and conditions defined in // file 'LICENSE.txt', which is part of this source code package. // *************************************************************** diff --git a/src/opt/gopt/setitem_gopt.cc b/src/opt/gopt/setitem_gopt.cc index 16867b5a..c777d2a0 100644 --- a/src/opt/gopt/setitem_gopt.cc +++ b/src/opt/gopt/setitem_gopt.cc @@ -48,7 +48,6 @@ static void setitem_inplace(SetitemOp* op) { } auto output = op->outputs().front(); output->share_with(input); - return; auto data = op->input(1); // if setitem requires type conversion, don't inplace diff --git a/src/pybind/py_var_tracer.cc b/src/pybind/py_var_tracer.cc index 36fe8352..e46b2afe 100644 --- a/src/pybind/py_var_tracer.cc +++ b/src/pybind/py_var_tracer.cc @@ -1,6 +1,9 @@ // *************************************************************** -// Copyright (c) 2021 Jittor. All Rights Reserved. -// Maintainers: Dun Liang . +// Copyright (c) 2021 Jittor. All Rights Reserved. +// Maintainers: +// Dun Liang . +// Guoye Yang <498731903@qq.com> +// // This file is subject to the terms and conditions defined in // file 'LICENSE.txt', which is part of this source code package. // *************************************************************** @@ -186,7 +189,7 @@ void TraceData::record_node(Node* node, bool record_stack) { NodeData data; data.id = node_data_cnt++; id_map[node] = data.id; - if (trace_py_var) { + if (!node->is_var() || trace_py_var>=3) { if (record_stack) { if (trace_grad_op) { auto iter = trace_data.id_map.find(trace_grad_op);