saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

delete model_copy to save memory allocated in forward call (#7832)

Browse source 
* delete model copy

* add flag

* address comments

* address flag comment

Co-authored-by: root <root@OrtTrainingDev0.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
harshithapv 2021-05-25 22:22:13 -07:00 committed by GitHub
parent 1c6b6f696e
commit 4fe59c8b29
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
orttraining/orttraining/python/training/ortmodule/_io.py
View file

@ -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