def _create_execution_agent(self):
"""Creates a TrainingAgent that can run the forward and backward graph on the training model"""
session_options, providers, provider_options = self._get_session_config(
)
fw_feed_names = [
input.name for input in self._optimized_onnx_model.graph.input
]
fw_outputs_device_info = []
for idx in range(len(self._graph_info.user_output_names)):
fw_outputs_device_info.append(
C.OrtDevice(get_ort_device_type(self._device.type),
C.OrtDevice.default_memory(),
_utils.get_device_index(self._device)))
bw_fetches_names = [
output.name for output in self._optimized_onnx_model.graph.output
]
bw_outputs_device_info = []
for idx in range(len(bw_fetches_names)):
bw_outputs_device_info.append(
C.OrtDevice(get_ort_device_type(self._device.type),
C.OrtDevice.default_memory(),
_utils.get_device_index(self._device)))
self._execution_agent = TrainingAgent(
self._optimized_onnx_model.SerializeToString(), fw_feed_names,
fw_outputs_device_info, bw_fetches_names, bw_outputs_device_info,
session_options, providers, provider_options)
def bind_output(self, name, device_type='cpu', device_id=0, element_type=None, shape=None, buffer_ptr=None):
'''
:param name: output name
:param device_type: e.g. cpu, cuda, cpu by default
:param device_id: device id, e.g. 0
:param element_type: output element type
:param shape: output shape
:param buffer_ptr: memory pointer to output data
'''
# Follow the `if` path when the user has not provided any pre-allocated buffer but still
# would like to bind an output to a specific device (e.g. cuda).
# Pre-allocating an output buffer may not be an option for the user as :
# (1) They may not want to use a custom allocator specific to the device they want to bind the output to,
# in which case ORT will allocate the memory for the user
# (2) The output has a dynamic shape and hence the size of the buffer may not be fixed across runs
if buffer_ptr is None:
self._iobinding.bind_output(name,
C.OrtDevice(get_ort_device_type(device_type), C.OrtDevice.default_memory(),
device_id))
else:
if element_type is None or shape is None:
raise ValueError("`element_type` and `shape` are to be provided if pre-allocated memory is provided")
self._iobinding.bind_output(name,
C.OrtDevice(get_ort_device_type(device_type), C.OrtDevice.default_memory(),
device_id),
element_type, shape, buffer_ptr)
def _create_execution_agent(self):
"""Creates a TrainingAgent that can run the forward and backward graph on the training model"""
session_options, providers, provider_options = self._get_session_config(
)
fw_feed_names = [
input.name
for input in self._onnx_models.optimized_model.graph.input
]
device_type = self._device if type(
self._device) is str else self._device.type.lower()
if device_type == "ort":
fw_outputs_device_info = [C.get_ort_device(self._device.index)] * (
len(self._graph_info.user_output_names) +
len(self._graph_info.frontier_node_arg_map))
else:
fw_outputs_device_info = [
C.OrtDevice(
get_ort_device_type(self._device.type, self._device.index),
C.OrtDevice.default_memory(),
_utils.get_device_index(self._device),
)
] * (len(self._graph_info.user_output_names) +
len(self._graph_info.frontier_node_arg_map))
bw_fetches_names = [
output.name
for output in self._onnx_models.optimized_model.graph.output
]
if device_type == "ort":
bw_outputs_device_info = [C.get_ort_device(self._device.index)
] * len(bw_fetches_names)
else:
bw_outputs_device_info = [
C.OrtDevice(
get_ort_device_type(self._device.type, self._device.index),
C.OrtDevice.default_memory(),
_utils.get_device_index(self._device),
)
] * len(bw_fetches_names)
local_device_rank = self._device.index if device_type == "ort" else _utils.get_device_index(
self._device)
self._execution_agent = TrainingAgent(
self._onnx_models.optimized_model.SerializeToString(),
fw_feed_names,
fw_outputs_device_info,
bw_fetches_names,
bw_outputs_device_info,
session_options,
providers,
provider_options,
local_device_rank,
)
def make(ort_device_name, device_id):
return OrtDevice(
C.OrtDevice(
get_ort_device_type(ort_device_name, device_id),
C.OrtDevice.default_memory(),
device_id,
))
def ortvalue_from_shape_and_type(shape=None,
element_type=None,
device_type="cpu",
device_id=0):
"""
Factory method to construct an OrtValue (which holds a Tensor) from given shape and element_type
:param shape: List of integers indicating the shape of the OrtValue
:param element_type: The data type of the elements in the OrtValue (numpy type)
:param device_type: e.g. cpu, cuda, cpu by default
:param device_id: device id, e.g. 0
"""
if shape is None or element_type is None:
raise ValueError(
"`element_type` and `shape` are to be provided if pre-allocated memory is provided"
)
return OrtValue(
C.OrtValue.ortvalue_from_shape_and_type(
shape,
element_type,
C.OrtDevice(
get_ort_device_type(device_type, device_id),
C.OrtDevice.default_memory(),
device_id,
),
))
def ortvalue_from_numpy(numpy_obj, device_type='cpu', device_id=0):
'''
Factory method to construct an OrtValue (which holds a Tensor) from a given Numpy object
A copy of the data in the Numpy object is held by the OrtValue only if the device is NOT cpu
:param numpy_obj: The Numpy object to construct the OrtValue from
:param device_type: e.g. cpu, cuda, cpu by default
:param device_id: device id, e.g. 0
'''
# Hold a reference to the numpy object (if device_type is 'cpu') as the OrtValue
# is backed directly by the data buffer of the numpy object and so the numpy object
# must be around until this OrtValue instance is around
return OrtValue(C.OrtValue.ortvalue_from_numpy(numpy_obj, C.OrtDevice(get_ort_device_type(device_type),
C.OrtDevice.default_memory(), device_id)), numpy_obj if device_type.lower() == 'cpu' else None)
def bind_input(self, name, device_type, device_id, element_type, shape, buffer_ptr):
'''
:param name: input name
:param device_type: e.g. cpu, cuda
:param device_id: device id, e.g. 0
:param element_type: input element type
:param shape: input shape
:param buffer_ptr: memory pointer to input data
'''
self._iobinding.bind_input(name,
C.OrtDevice(get_ort_device_type(device_type), C.OrtDevice.default_memory(),
device_id),
element_type, shape, buffer_ptr)
def bind_output(self,
name,
device_type='cpu',
device_id=0,
element_type=None,
shape=None,
buffer_ptr=None):
'''
:param name: output name
:param device_type: e.g. CPU, CUDA, CPU by default
:param device_id: device id, e.g. 0
:param element_type: output element type
:param shape: output shape
:param buffer_ptr: memory pointer to output data
'''
if device_type == 'cpu':
self._iobinding.bind_output(name)
else:
self._iobinding.bind_output(
name,
C.OrtDevice(get_ort_device_type(device_type),
C.OrtDevice.default_memory(), device_id),
element_type, shape, buffer_ptr)
def bind_output(self, name, device_type, device_id, element_type, shape, buffer_ptr):
self._iobinding.bind_output(name,
C.OrtDevice(get_ort_device_type(device_type), C.OrtDevice.default_memory(),
device_id),
element_type, shape, buffer_ptr)
