From 4fe59c8b29613fcd9b8fb6125a05358164f81bb5 Mon Sep 17 00:00:00 2001 From: harshithapv <54084812+harshithapv@users.noreply.github.com> Date: Tue, 25 May 2021 22:22:13 -0700 Subject: [PATCH] delete model_copy to save memory allocated in forward call (#7832) * delete model copy * add flag * address comments * address flag comment Co-authored-by: root --- .../orttraining/python/training/ortmodule/_io.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule/_io.py b/orttraining/orttraining/python/training/ortmodule/_io.py index b7cb3e79a8..8aaa9e908c 100644 --- a/orttraining/orttraining/python/training/ortmodule/_io.py +++ b/orttraining/orttraining/python/training/ortmodule/_io.py @@ -8,7 +8,7 @@ import copy import inspect import torch import warnings - +import gc class _OutputIdentityOp(torch.autograd.Function): '''Internal class used to prepend Identity ops in model's outputs @@ -460,12 +460,14 @@ def parse_outputs_for_onnx_export_and_extract_schema(module, inputs, kwargs): module.eval() output_names = None output_dynamic_axes = None + is_deepcopy = False with torch.no_grad(): # Deepcopy inputs, since input values may change after model run. sample_inputs_copy, sample_kwargs_copy = deepcopy_model_input(*inputs, **kwargs) try: # Deepcopy model, in case model is stateful and changes after model run. model_copy = copy.deepcopy(module) + is_deepcopy = True except Exception: model_copy = module warnings.warn("This model cannot be deep copied (or pickled), " @@ -478,6 +480,9 @@ def parse_outputs_for_onnx_export_and_extract_schema(module, inputs, kwargs): output_names, output_dynamic_axes = _parse_outputs_and_extract_names_and_dynamic_axes(sample_outputs) if is_train_mode: module.train() - + output_schema = _extract_schema(sample_outputs) + if is_deepcopy: + del model_copy + gc.collect() # Return output names, output dynamic axes and output schema - return output_names, output_dynamic_axes, _extract_schema(sample_outputs) + return output_names, output_dynamic_axes, output_schema