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

Enable symbolic shape inference in ORTModule (#7282)

Browse source 
Co-authored-by: Ethan Tao <ettao@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
ytaous 2021-04-08 09:47:09 -07:00 committed by GitHub
parent d272c8434d
commit e14b291ce7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
orttraining/orttraining/python/training/_ortmodule_graph_execution_manager.py
View file

@ -9,6 +9,8 @@ from . import _ortmodule_logger as _logger
from onnxruntime.capi.onnxruntime_inference_collection import OrtValue
from onnxruntime.capi import _pybind_state as C
from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
from abc import ABC, abstractmethod
import io
import inspect
@ -90,6 +92,8 @@ class GraphExecutionManager(ABC):
# default execution order is priority-based for both dynamic/static shape input for now
# if we observe benefit of static shape, we can expose this flag to user
self._use_static_shape = False
# flag to enable symbolic shape inference for dynamic shape inputs to improve performance
self._run_symbolic_shape_infer = False
self._input_names_require_grad = None
self._module_output_schema = None
@ -184,6 +188,9 @@ class GraphExecutionManager(ABC):
self._set_device_from_module()
self._onnx_model = self._get_exported_model(*inputs, **kwargs)
if self._run_symbolic_shape_infer:
self._onnx_model = SymbolicShapeInference.infer_shapes(self._onnx_model, auto_merge=True, guess_output_rank=True)
return True
def _get_exported_model(self, *inputs, **kwargs):