From 92bebe65a9de22130b46e35d20aac92251bc916f Mon Sep 17 00:00:00 2001
From: kingxian <jinxiaoxian@huawei.com>
Date: Wed, 2 Sep 2020 17:13:05 +0800
Subject: [PATCH] cell init performance opt

---
 mindspore/nn/cell.py                | 6 +++++-
 tests/st/ops/cpu/test_maxpool_op.py | 4 ++--
 tests/ut/cpp/runtest.sh             | 1 +
 tests/ut/python/runtest.sh          | 1 +
 4 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/mindspore/nn/cell.py b/mindspore/nn/cell.py
index e50e93e8143..ed6eac40363 100755
--- a/mindspore/nn/cell.py
+++ b/mindspore/nn/cell.py
@@ -16,6 +16,7 @@
 import inspect
 import time
 import gc
+import os
 from collections import OrderedDict
 import numpy
 from mindspore import log as logger
@@ -73,8 +74,11 @@ class Cell:
         self._parameter_layout_dict = {}
         self._create_time = int(time.time() * 1e9)
         init_backend()
+
         # call gc to release GE session resources used by non-used cell objects
-        gc.collect()
+        if os.getenv('GC_COLLECT_IN_CELL') == '1':
+            gc.collect()
+
         self._construct_inputs_num = 0
         self._construct_inputs_names = []
         self._auto_parallel_mode = False
diff --git a/tests/st/ops/cpu/test_maxpool_op.py b/tests/st/ops/cpu/test_maxpool_op.py
index d13e1cecb9f..923607da02b 100644
--- a/tests/st/ops/cpu/test_maxpool_op.py
+++ b/tests/st/ops/cpu/test_maxpool_op.py
@@ -35,10 +35,10 @@ class Net_Pool(nn.Cell):
 class Net_Pool2(nn.Cell):
     def __init__(self):
         super(Net_Pool2, self).__init__()
-        self.maxpool_fun = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode="SAME")
+        self.maxpool_fun2 = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode="SAME")
 
     def construct(self, x):
-        return self.maxpool_fun(x)
+        return self.maxpool_fun2(x)
 
 
 @pytest.mark.level0
diff --git a/tests/ut/cpp/runtest.sh b/tests/ut/cpp/runtest.sh
index a0503a0e94e..db720e0deae 100755
--- a/tests/ut/cpp/runtest.sh
+++ b/tests/ut/cpp/runtest.sh
@@ -29,6 +29,7 @@ cd ${BUILD_PATH}/mindspore/tests/ut/cpp
 export LD_LIBRARY_PATH=${BUILD_PATH}/mindspore/googletest/googlemock/gtest:${PROJECT_PATH}/mindspore:${PROJECT_PATH}/mindspore/lib:$LD_LIBRARY_PATH
 export PYTHONPATH=${PROJECT_PATH}/tests/ut/cpp/python_input:$PYTHONPATH:${PROJECT_PATH}
 export GLOG_v=2
+export GC_COLLECT_IN_CELL=1
 
 ## prepare data for dataset & mindrecord
 cp -fr $PROJECT_PATH/tests/ut/data ${PROJECT_PATH}/build/mindspore/tests/ut/cpp/
diff --git a/tests/ut/python/runtest.sh b/tests/ut/python/runtest.sh
index 1a687b9b35d..10a6fafbe15 100755
--- a/tests/ut/python/runtest.sh
+++ b/tests/ut/python/runtest.sh
@@ -27,6 +27,7 @@ fi
 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${BUILD_PATH}/third_party/gtest/lib
 export PYTHONPATH=$PYTHONPATH:${PROJECT_PATH}:${PROJECT_PATH}/tests/ut/cpp/python_input:${PROJECT_PATH}/tests/ut/python
 echo "export PYTHONPATH=$PYTHONPATH"
+export GC_COLLECT_IN_CELL=1
 
 if [ $# -eq 1 ]  &&  ([ "$1" == "stage1" ] || [ "$1" == "stage2" ] || [ "$1" == "stage3" ]); then
     if [ $1 == "stage1" ]; then