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

Cache initializers and avoid device check ot end of forward (#7905)

Browse source
This commit is contained in:
baijumeswani 2021-06-02 10:07:49 -07:00 committed by GitHub
parent 9946e6f7df
commit 271a343024
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 8 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
View file

@ -54,6 +54,7 @@ class GraphExecutionManager(ABC):
self._graph_info = None
self._graph_initializer_names = None
self._graph_initializer_names_to_train = None
self._graph_initializers = None
# TrainingAgent or InferenceAgent
self._execution_agent = None
@ -320,3 +321,8 @@ class GraphExecutionManager(ABC):
# a set (unordered_set in the backend) that does not require a copy on each reference.
self._graph_initializer_names = set(initializer_names)
self._graph_initializer_names_to_train = set(initializer_names_to_train)
# Initializers can be cached and used since they are expected not to be re-instantiated
# between forward calls.
self._graph_initializers = [param for name, param in self._flattened_module.named_parameters()
if name in self._graph_initializer_names]

3
orttraining/orttraining/python/training/ortmodule/_inference_manager.py
View file

@ -89,8 +89,7 @@ class InferenceManager(GraphExecutionManager):
self._optimized_onnx_model,
self._device,
*_io._combine_input_buffers_initializers(
[param for name, param in self._flattened_module.named_parameters()
if name in self._graph_initializer_names],
self._graph_initializers,
self._graph_info.user_input_names,
self._input_info,
self._flattened_module.named_buffers(),

6
orttraining/orttraining/python/training/ortmodule/_training_manager.py
View file

@ -47,9 +47,6 @@ class TrainingManager(GraphExecutionManager):
execution_session.run_forward(forward_inputs, forward_outputs, state)
user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output) for forward_output in forward_outputs)
# Assert that the outputs and model device match
_utils._check_same_device(device, "Output argument from forward", *user_outputs)
output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
run_info = RunStateInfo(state, output_info)
# Return user outputs and forward run information
@ -192,8 +189,7 @@ class TrainingManager(GraphExecutionManager):
return _io.unflatten_user_output(self._module_output_schema,
_ORTModuleFunction.apply(
*_io._combine_input_buffers_initializers(
[param for name, param in self._flattened_module.named_parameters()
if name in self._graph_initializer_names],
self._graph_initializers,
self._graph_info.user_input_names,
self._input_info,
self._flattened_module.named_buffers(),