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

Avoid duplicated calls of postprocess in training frontend (#4579)

Browse source
This commit is contained in:
Bowen Bao 2020-08-07 21:34:11 -07:00 committed by GitHub
parent 77c69a0325
commit abbb7f6f5c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 17 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
orttraining/orttraining/python/ort_trainer.py
View file

@ -278,7 +278,7 @@ def wrap_for_input_match(model, loss_fn, input_names):
return model
def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, opset_version=DEFAULT_OPSET_VERSION, _enable_internal_postprocess=True):
def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, opset_version=DEFAULT_OPSET_VERSION):
# example: {input0:{0:'batch'}, input1:{0:'batch'}}
dynamic_axes = {}
for input in model_desc.inputs_:
@ -367,9 +367,6 @@ def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, op
"Initializer names do not match between PyTorch model and ONNX model, " \
"please report a bug to ONNX Runtime."
if _enable_internal_postprocess:
onnx_model = postprocess.run_postprocess(onnx_model)
return onnx_model
def create_ort_training_session_with_optimizer(model, device, training_optimizer_name, lr_params_feed_name,
@ -609,6 +606,9 @@ class ORTTrainer():
self.torch_model_ = None
self.onnx_model_ = None
self._enable_internal_postprocess = _enable_internal_postprocess
self._extra_postprocess = _extra_postprocess
if isinstance(model, torch.nn.Module):
self.torch_model_ = model
self.loss_fn_ = loss_fn
@ -619,6 +619,12 @@ class ORTTrainer():
# TODO: accept loss_fn as an onnx model. build self.onnx_model_ with model and loss_fn
self.loss_fn_ = None
if self._enable_internal_postprocess:
postprocess.run_postprocess(self.onnx_model_)
if self._extra_postprocess:
self._extra_postprocess(self.onnx_model_)
self.model_desc_ = model_desc
self.input_desc_with_lr = [*self.model_desc_.inputs_, learning_rate_description]
@ -640,7 +646,6 @@ class ORTTrainer():
# we use self.global_step_ to count optimizations being performed.
# it is used to calculate learning rate if self.get_lr_this_step_ is provided.
self.global_step_ = global_step
self._extra_postprocess = _extra_postprocess
self.get_lr_this_step_ = get_lr_this_step
self.loss_scaler_ = loss_scaler
@ -655,7 +660,6 @@ class ORTTrainer():
self.frozen_weights_ = frozen_weights
self.opset_version_ = _opset_version
self.state_dict_ = None
self._enable_internal_postprocess = _enable_internal_postprocess
self._use_deterministic_compute = _use_deterministic_compute
self.use_invertible_layernorm_grad = use_invertible_layernorm_grad
@ -669,12 +673,6 @@ class ORTTrainer():
if self.onnx_model_ is None:
return
if self._enable_internal_postprocess:
self._onnx_model_ = postprocess.run_postprocess(self.onnx_model_)
if self._extra_postprocess:
self._extra_postprocess(self.onnx_model_)
self._verify_fully_optimized_model(self.onnx_model_)
self.session, self.train_io_binding, self.eval_io_binding, self.output_name, _, self.output_types = \
create_ort_training_session_with_optimizer(
@ -732,7 +730,13 @@ class ORTTrainer():
torch_buffers = list(dict(self.torch_model_.named_buffers()).keys())
self.frozen_weights_ = self.frozen_weights_ + torch_buffers
self.onnx_model_ = convert_model_loss_fn_to_onnx(
self.torch_model_, self.loss_fn_, self.model_desc_, torch.device('cpu'), inputs, opset_version=self.opset_version_, _enable_internal_postprocess=self._enable_internal_postprocess)
self.torch_model_, self.loss_fn_, self.model_desc_, torch.device('cpu'), inputs, opset_version=self.opset_version_)
if self._enable_internal_postprocess:
postprocess.run_postprocess(self.onnx_model_)
if self._extra_postprocess:
self._extra_postprocess(self.onnx_model_)
self._init_session()