Add memory_optimize in Transformer

fluid/neural_machine_translation/transformer/train.py

@@ -107,6 +107,11 @@ def parse_args():
         default=False,
         help="The flag indicating whether to run the task "
         "for continuous evaluation.")
+    parser.add_argument(
+        "--use_mem_opt",
+        type=ast.literal_eval,
+        default=True,
+        help="The flag indicating whether to use memory optimization.")
 
     args = parser.parse_args()
     # Append args related to dict
@@ -515,6 +520,9 @@ def train(args):
             epsilon=TrainTaskConfig.eps)
         optimizer.minimize(sum_cost)
 
+    if args.use_mem_opt:
+        fluid.memory_optimize(fluid.default_main_program())
+
     if args.local:
         print("local start_up:")
         train_loop(exe,