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

move some process out of training step (#11150)

Browse source
This commit is contained in:
Vincent Wang 2022-04-08 17:30:11 +08:00 committed by GitHub
parent 4c37f15c1b
commit bcc62e0cbf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 21 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -246,6 +246,26 @@ class GraphExecutionManager(GraphExecutionInterface):
self._graph_info = self._graph_builder.get_graph_info()
# Map each input/initializer to its gradient index in the graph output, or -1 is gradient is not required.
self._gradient_map = []
num_user_input_grads = len(self._input_info.require_grad_names)
require_grad_names_set = set(self._input_info.require_grad_names)
require_grad_names_index = 0
for input_name in self._graph_info.user_input_names:
if input_name in require_grad_names_set:
self._gradient_map.append(require_grad_names_index)
require_grad_names_index += 1
else:
self._gradient_map.append(-1)
initializer_index = num_user_input_grads
for initializer_name in self._graph_info.initializer_names:
if initializer_name in self._graph_initializer_names_to_train:
self._gradient_map.append(initializer_index)
initializer_index += 1
else:
self._gradient_map.append(-1)
def _get_session_config(self):
"""Creates and returns the session configuration to be used for the ExecutionAgent"""

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

@ -164,37 +164,11 @@ class TrainingManager(GraphExecutionManager):
# Destroy the state immediately (as opposed to be at the mercy of garbage collector) so it does not
# affect peak memory usage in a subsequent graph run.
del ctx.run_info.state
num_user_input_grads = len(self._input_info.require_grad_names)
results = []
require_grad_names_set = set(self._input_info.require_grad_names)
require_grad_names_index = 0
# Fast version: all backward_outputs are converted first.
# This version only works if backward_outputs is an OrtValueVector.
transfered_backward_outputs = _utils._ortvalues_to_torch_tensor(backward_outputs, self._device)
# Return input and initializer gradients
for input_name in self._graph_info.user_input_names:
# Append to the results the backward output for each input that required grad
if input_name in require_grad_names_set:
results.append(transfered_backward_outputs[require_grad_names_index])
require_grad_names_index += 1
else:
# input_name is not found in the self._input_info.require_grad_names list
# Append None to results for each input that did not require grad
results.append(None)
# Append gradients of initializer to results
# Go over each initializer, check if it required grad and append to results accordingly
initializer_index = num_user_input_grads
for initializer_name in self._graph_info.initializer_names:
if initializer_name in self._graph_initializer_names_to_train:
results.append(transfered_backward_outputs[initializer_index])
initializer_index += 1
else:
results.append(None)
return tuple(results)
return tuple(transfered_backward_outputs[idx] if idx != -1 else None for idx in self._gradient_map)
return _ORTModuleFunction