add mem_opt argument

dzhwinter · Mar 28, 2018 · 5ededc4 · 5ededc4
1 parent daea299
commit 5ededc4
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/fluid/SE-ResNeXt-152/train_parallel_do.py b/fluid/SE-ResNeXt-152/train_parallel_do.py
@@ -26,6 +26,11 @@
 def parse_args():
     parser = argparse.ArgumentParser('SE-ResNeXt-152 parallel profile.')
     parser.add_argument('--per_gpu_batch_size', type=int, default=12, help='')
+    parser.add_argument(
+        '--use_mem_opt',
+        type=distutils.util.strtobool,
+        default=True,
+        help='use memory optimize')
     parser.add_argument(
         '--skip_first_steps',
         type=int,
@@ -212,7 +217,8 @@ def train():
         regularization=fluid.regularizer.L2Decay(1e-4))
     opts = optimizer.minimize(avg_cost)
 
-    fluid.memory_optimize(fluid.default_main_program())
+    if args.use_mem_opt:
+        fluid.memory_optimize(fluid.default_main_program())
 
     place = fluid.CUDAPlace(0)
     # place = fluid.CPUPlace()