Support OrtValue binding in Python to enable interesting IOBinding scenarios in Python (#5248)

This commit is contained in:
Hariharan Seshadri 2020-10-06 21:14:41 -07:00 committed by GitHub
parent 0122e890d9
commit 6f54113a1b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 845 additions and 288 deletions

View file

@ -9,6 +9,34 @@ in *ONNX Runtime*.
.. contents::
:local:
OrtValue
=========
*ONNX Runtime* works with native Python data structures which are mapped into ONNX ONNX data formats :
Numpy arrays (tensors), dictionaries (maps), and a list of Numpy arrays (sequences).
The data backing these are on CPU.
*ONNX Runtime* supports a custom data structure that supports all ONNX data formats that allows users
to place the data backing these on a device, for example, on a CUDA supported device. This allows for
interesting *IOBinding* scenarios (discussed below). In addition, *ONNX Runtime* supports directly
working with *OrtValue*(s) while inferencing a model if provided as part of the input feed.
Below is an example showing creation of an *OrtValue* from a Numpy array while placing its backing memory
on a CUDA device:
.. code-block:: python
#X is numpy array on cpu, create an OrtValue and place it on cuda device id = 0
ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(X, 'cuda', 0)
ortvalue.device_name() # 'cuda'
ortvalue.shape() # shape of the numpy array X
ortvalue.data_type() # 'tensor(float)'
ortvalue.is_tensor() # 'True'
np.array_equal(ortvalue.numpy(), X) # 'True'
#ortvalue can be provided as part of the input feed to a model
ses = onnxruntime.InferenceSession('model.onnx')
res = sess.run(["Y"], {"X": ortvalue})
IOBinding
=========
@ -24,7 +52,7 @@ Here are scenarios to use this feature.
Scenario 1:
A graph is executed on a deivce other than CPU, for instance CUDA. Users can
A graph is executed on a device other than CPU, for instance CUDA. Users can
use IOBinding to put input on CUDA as the follows.
.. code-block:: python
@ -32,6 +60,7 @@ use IOBinding to put input on CUDA as the follows.
#X is numpy array on cpu
session = onnxruntime.InferenceSession('model.onnx')
io_binding = session.io_binding()
# OnnxRuntime will copy the data over to the CUDA device if 'input' is consumed by nodes on the CUDA device
io_binding.bind_cpu_input('input', X)
io_binding.bind_output('output')
session.run_with_iobinding(io_binding)
@ -39,27 +68,68 @@ use IOBinding to put input on CUDA as the follows.
Scenario 2:
The input data is on a device, users direclty use the input. The output data is on CPU.
The input data is on a device, users directly use the input. The output data is on CPU.
.. code-block:: python
#X is numpy array on cpu
X_ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(X, 'cuda', 0)
session = onnxruntime.InferenceSession('model.onnx')
io_binding = session.io_binding()
io_binding.bind_input(name='input', device_type=X.device.type, device_id=0, element_type=np.float32, shape=list(X.size()), buffer_ptr=X.data_ptr())
io_binding.bind_input(name='input', device_type=X_ortvalue.device_name(), device_id=0, element_type=np.float32, shape=X_ortvalue.shape(), buffer_ptr=X_ortvalue.data_ptr())
io_binding.bind_output('output')
session.run_with_iobinding(io_binding)
Y = io_binding.copy_outputs_to_cpu()[0]
Scenario 3:
The input data on a dveice, users directly use the input and also place output on the device:
The input data and output data are both on a device, users directly use the input and also place output on the device.
.. code-block:: python
#X is numpy array on cpu
X_ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(X, 'cuda', 0)
Y_ortvalue = onnxruntime.OrtValue.ortvalue_from_shape_and_type([3, 2], np.float32, 'cuda', 0) # Change the shape to the actual shape of the output being bound
session = onnxruntime.InferenceSession('model.onnx')
io_binding = session.io_binding()
io_binding.bind_input(name='input', device_type=X.device.type, device_id=0, element_type=np.float32, shape=list(X.size()), buffer_ptr=X.data_ptr())
io_binding.bind_output(name='output', device_type=Y.device.type, device_id=0, element_type=np.float32, shape=list(Y.size()), buffer_ptr=Y.data_ptr())
io_binding.bind_input(name='input', device_type=X_ortvalue.device_name(), device_id=0, element_type=np.float32, shape=X_ortvalue.shape(), buffer_ptr=X_ortvalue.data_ptr())
io_binding.bind_output(name='output', device_type=Y_ortvalue.device_name(), device_id=0, element_type=np.float32, shape=Y_ortvalue.shape(), buffer_ptr=Y_ortvalue.data_ptr())
session.run_with_iobinding(io_binding)
Scenario 4:
Users can request *ONNX Runtime* to allocate an output on a device. This is particularly useful for dynamic shaped outputs.
Users can use the `get_outputs()` API to get the *OrtValue*(s) corresponding to the allocated output(s).
Users can thus consume the *ONNX Runtime* allocated memory for the output as an *OrtValue*.
.. code-block:: python
#X is numpy array on cpu
X_ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(X, 'cuda', 0)
session = onnxruntime.InferenceSession('model.onnx')
io_binding = session.io_binding()
io_binding.bind_input(name='input', device_type=X_ortvalue.device_name(), device_id=0, element_type=np.float32, shape=X_ortvalue.shape(), buffer_ptr=X_ortvalue.data_ptr())
#Request ONNX Runtime to bind and allocate memory on CUDA for 'output'
io_binding.bind_output('output', 'cuda')
session.run_with_iobinding(io_binding)
# The following call returns an OrtValue which has data allocated by ONNX Runtime on CUDA
ort_output = io_binding.get_outputs()[0]
Scenario 5:
Users can bind *OrtValue*(s) directly.
.. code-block:: python
#X is numpy array on cpu
#X is numpy array on cpu
X_ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(X, 'cuda', 0)
Y_ortvalue = onnxruntime.OrtValue.ortvalue_from_shape_and_type([3, 2], np.float32, 'cuda', 0) # Change the shape to the actual shape of the output being bound
session = onnxruntime.InferenceSession('model.onnx')
io_binding = session.io_binding()
io_binding.bind_ortvalue_input('input', X_ortvalue)
io_binding.bind_ortvalue_output('output', Y_ortvalue)
session.run_with_iobinding(io_binding)
Device

View file

@ -19,7 +19,7 @@ try:
except ImportError:
pass
from onnxruntime.capi.session import InferenceSession, IOBinding
from onnxruntime.capi.onnxruntime_inference_collection import InferenceSession, IOBinding, OrtValue
from onnxruntime.capi import onnxruntime_validation
from onnxruntime.capi.training import * # noqa: F403

View file

@ -8,6 +8,7 @@ from onnxruntime.capi import _pybind_state as C
def get_ort_device_type(device):
device = device.lower()
if device == 'cuda':
return C.OrtDevice.cuda()
elif device == 'cpu':
@ -253,6 +254,7 @@ class IOBinding:
'''
def __init__(self, session):
self._iobinding = C.SessionIOBinding(session._sess)
self._numpy_obj_references = []
def bind_cpu_input(self, name, arr_on_cpu):
'''
@ -260,12 +262,16 @@ class IOBinding:
:param name: input name
:param arr_on_cpu: input values as a python array on CPU
'''
# Hold a reference to the numpy object as the bound OrtValue is backed
# directly by the data buffer of the numpy object and so the numpy object
# must be around until this IOBinding instance is around
self._numpy_obj_references.append(arr_on_cpu)
self._iobinding.bind_input(name, arr_on_cpu)
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_type: e.g. cpu, cuda
:param device_id: device id, e.g. 0
:param element_type: input element type
:param shape: input shape
@ -276,10 +282,17 @@ class IOBinding:
device_id),
element_type, shape, buffer_ptr)
def bind_ortvalue_input(self, name, ortvalue):
'''
:param name: input name
:param ortvalue: OrtValue instance to bind
'''
self._iobinding.bind_ortvalue_input(name, ortvalue._ortvalue)
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_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
@ -304,6 +317,25 @@ class IOBinding:
device_id),
element_type, shape, buffer_ptr)
def bind_ortvalue_output(self, name, ortvalue):
'''
:param name: output name
:param ortvalue: OrtValue instance to bind
'''
self._iobinding.bind_ortvalue_output(name, ortvalue._ortvalue)
def get_outputs(self):
'''
Returns the output OrtValues from the Run() that preceded the call.
The data buffer of the obtained OrtValues may not reside on CPU memory
'''
returned_ortvalues = []
for ortvalue in self._iobinding.get_outputs():
returned_ortvalues.append(OrtValue(ortvalue))
return returned_ortvalues
def copy_outputs_to_cpu(self):
'''Copy output contents to CPU (if on another device). No-op if already on the CPU.'''
return self._iobinding.copy_outputs_to_cpu()
@ -313,3 +345,88 @@ class IOBinding:
def clear_binding_outputs(self):
self._iobinding.clear_binding_outputs()
class OrtValue:
'''
A data structure that supports all ONNX data formats (tensors and non-tensors) that allows users
to place the data backing these on a device, for example, on a CUDA supported device.
This class provides APIs to construct and deal with OrtValues.
'''
def __init__(self, ortvalue, numpy_obj=None):
if isinstance(ortvalue, C.OrtValue):
self._ortvalue = ortvalue
# Hold a ref count to the numpy object if the OrtValue is backed directly
# by its data buffer so that it isn't destroyed when the OrtValue is in use
self._numpy_obj = numpy_obj
else:
# An end user won't hit this error
raise ValueError("`Provided ortvalue` needs to be of type " +
"`onnxruntime.capi.onnxruntime_pybind11_state.OrtValue`")
@staticmethod
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)
@staticmethod
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), C.OrtDevice.default_memory(), device_id)))
def data_ptr(self):
'''
Returns the address of the first element in the OrtValue's data buffer
'''
return self._ortvalue.data_ptr()
def device_name(self):
'''
Returns the name of the device where the OrtValue's data buffer resides e.g. cpu, cuda
'''
return self._ortvalue.device_name().lower()
def shape(self):
'''
Returns the shape of the data in the OrtValue
'''
return self._ortvalue.shape()
def data_type(self):
'''
Returns the data type of the data in the OrtValue
'''
return self._ortvalue.data_type()
def is_tensor(self):
'''
Returns True if the OrtValue is a Tensor, else returns False
'''
return self._ortvalue.is_tensor()
def numpy(self):
'''
Returns a Numpy object from the OrtValue.
Valid only for OrtValues holding Tensors. Throws for OrtValues holding non-Tensors.
'''
return self._ortvalue.numpy()

View file

@ -23,6 +23,38 @@ namespace python {
namespace py = pybind11;
using namespace onnxruntime::logging;
static bool PyObjectCheck_NumpyArray(PyObject* o) {
return PyObject_HasAttrString(o, "__array_finalize__");
}
bool IsNumericNumpyType(int npy_type) {
return npy_type < NPY_OBJECT || npy_type == NPY_HALF;
}
bool IsNumericNumpyArray(py::object& py_object) {
if (PyObjectCheck_NumpyArray(py_object.ptr())) {
int npy_type = PyArray_TYPE(reinterpret_cast<PyArrayObject*>(py_object.ptr()));
return IsNumericNumpyType(npy_type);
}
return false;
}
static TensorShape GetArrayShape(PyArrayObject* pyObject) {
int ndim = PyArray_NDIM(pyObject);
const npy_intp* npy_dims = PyArray_DIMS(pyObject);
std::vector<int64_t> dims(ndim);
for (int i = 0; i < ndim; ++i) {
dims[i] = npy_dims[i];
}
TensorShape shape(std::move(dims));
return shape;
}
void CpuToCpuMemCpy(void* dst, const void* src, size_t num_bytes) {
memcpy(dst, src, num_bytes);
}
int OnnxRuntimeTensorToNumpyType(const DataTypeImpl* tensor_type) {
static std::map<MLDataType, int> type_map{
{DataTypeImpl::GetType<bool>(), NPY_BOOL},
@ -118,17 +150,6 @@ const DataTypeImpl* NumpyToOnnxRuntimeTensorType(int numpy_type) {
}
}
TensorShape GetArrayShape(PyArrayObject* pyObject) {
int ndim = PyArray_NDIM(pyObject);
const npy_intp* npy_dims = PyArray_DIMS(pyObject);
std::vector<int64_t> dims(ndim);
for (int i = 0; i < ndim; ++i) {
dims[i] = npy_dims[i];
}
TensorShape shape(std::move(dims));
return shape;
}
// This is a one time use, ad-hoc allocator that allows Tensors to take ownership of
// python array objects and use the underlying memory directly and
// properly deallocated them when they are done.
@ -190,13 +211,11 @@ class OrtPybindSingleUseAllocator : public IAllocator {
using OrtPybindSingleUseAllocatorPtr = std::shared_ptr<OrtPybindSingleUseAllocator>;
bool PyObjectCheck_Array(PyObject* o) {
return PyObject_HasAttrString(o, "__array_finalize__");
}
// Expects p_tensor properly created
// Does not manage darray life-cycle
void CopyDataToTensor(PyArrayObject* darray, int npy_type, std::unique_ptr<Tensor>& p_tensor) {
static void CopyDataToTensor(PyArrayObject* darray, int npy_type, std::unique_ptr<Tensor>& p_tensor,
MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) {
const auto total_items = p_tensor->Shape().Size();
if (npy_type == NPY_UNICODE) {
// Copy string data which needs to be done after Tensor is allocated.
@ -251,11 +270,16 @@ void CopyDataToTensor(PyArrayObject* darray, int npy_type, std::unique_ptr<Tenso
if (!IAllocator::CalcMemSizeForArray(p_tensor->DataType()->Size(), p_tensor->Shape().Size(), &len)) {
throw std::runtime_error("length overflow");
}
memcpy(buffer, PyArray_DATA(darray), len);
mem_cpy_to_device(buffer, PyArray_DATA(darray), len);
}
}
std::unique_ptr<Tensor> CreateTensor(const AllocatorPtr& alloc, const std::string& name_input, PyArrayObject* pyObject) {
// Setting `use_numpy_data_memory` to `true` will ensure that the underlying numpy array buffer is directly used
// as the backing data buffer for the ORT Tensor where applicable (for numeric tensors)
// The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope
static std::unique_ptr<Tensor> CreateTensor(const AllocatorPtr& alloc, const std::string& name_input,
PyArrayObject* pyObject, bool use_numpy_data_memory = true,
MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) {
PyArrayObject* darray = PyArray_GETCONTIGUOUS(pyObject);
ORT_ENFORCE(darray != nullptr, "The object must be a contiguous array for input '", name_input, "'.");
@ -265,8 +289,7 @@ std::unique_ptr<Tensor> CreateTensor(const AllocatorPtr& alloc, const std::strin
const int npy_type = PyArray_TYPE(darray);
TensorShape shape = GetArrayShape(darray);
auto element_type = NumpyToOnnxRuntimeTensorType(npy_type);
if (npy_type != NPY_UNICODE && npy_type != NPY_STRING &&
npy_type != NPY_VOID && npy_type != NPY_OBJECT) {
if (IsNumericNumpyType(npy_type) && use_numpy_data_memory) {
if (pyObject == darray) {
// Use the memory of numpy array directly. The ownership belongs to the calling
// python code. In this case, the incoming pyObject must itself be contiguous (pyObject == darray).
@ -280,15 +303,15 @@ std::unique_ptr<Tensor> CreateTensor(const AllocatorPtr& alloc, const std::strin
}
} else {
p_tensor = onnxruntime::make_unique<Tensor>(element_type, shape, alloc);
CopyDataToTensor(darray, npy_type, p_tensor);
CopyDataToTensor(darray, npy_type, p_tensor, mem_cpy_to_device);
}
return p_tensor;
}
bool CheckIfInputIsSequenceType(const std::string& name_input,
const InputDefList* input_def_list,
/*out*/ onnx::TypeProto& type_proto) {
static bool CheckIfInputIsSequenceType(const std::string& name_input,
const InputDefList* input_def_list,
/*out*/ onnx::TypeProto& type_proto) {
// get sequence type from the model
const auto& def_list = *input_def_list;
auto ret_it = std::find_if(std::begin(def_list), std::end(def_list),
@ -306,8 +329,8 @@ bool CheckIfInputIsSequenceType(const std::string& name_input,
return type_proto.has_sequence_type();
}
void CreateSequenceOfTensors(AllocatorPtr alloc, const std::string& name_input,
const InputDefList* input_def_list, PyObject* pylist_obj, OrtValue* p_mlvalue) {
static void CreateSequenceOfTensors(AllocatorPtr alloc, const std::string& name_input,
const InputDefList* input_def_list, PyObject* pylist_obj, OrtValue* p_mlvalue) {
onnx::TypeProto type_proto;
if (!CheckIfInputIsSequenceType(name_input, input_def_list, type_proto)) {
throw std::runtime_error("Input is not of sequence type");
@ -320,7 +343,7 @@ void CreateSequenceOfTensors(AllocatorPtr alloc, const std::string& name_input,
tensors.resize(list_size);
for (Py_ssize_t i = 0; i < list_size; ++i) {
auto* py_obj = PyList_GetItem(pylist_obj, i);
if (!PyObjectCheck_Array(py_obj)) {
if (!PyObjectCheck_NumpyArray(py_obj)) {
throw std::runtime_error("CreateSequenceOfTensors: Input is not a tensor");
}
auto p_tensor = CreateTensor(alloc, name_input, reinterpret_cast<PyArrayObject*>(py_obj));
@ -339,9 +362,12 @@ void CreateSequenceOfTensors(AllocatorPtr alloc, const std::string& name_input,
ml_tensor_sequence->GetDeleteFunc());
}
void CreateTensorMLValue(const AllocatorPtr& alloc, const std::string& name_input, PyArrayObject* pyObject,
OrtValue* p_mlvalue) {
auto p_tensor = CreateTensor(alloc, name_input, pyObject);
// Setting `use_numpy_data_memory` to `true` will ensure that the underlying numpy array buffer is directly used
// as the backing data buffer for the ORT Tensor where applicable (for numeric tensors)
// The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope
static void CreateTensorMLValue(const AllocatorPtr& alloc, const std::string& name_input, PyArrayObject* pyObject,
OrtValue* p_mlvalue, bool use_numpy_data_memory = true, MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) {
auto p_tensor = CreateTensor(alloc, name_input, pyObject, use_numpy_data_memory, mem_cpy_to_device);
auto ml_tensor = DataTypeImpl::GetType<Tensor>();
p_mlvalue->Init(p_tensor.release(),
@ -351,7 +377,7 @@ void CreateTensorMLValue(const AllocatorPtr& alloc, const std::string& name_inpu
// This function will create a Tensor that owns the python array memory. This is done to properly
// release python arrays allocated within the pybind code.
void CreateTensorMLValueOwned(const OrtPybindSingleUseAllocatorPtr& pybind_alloc, const AllocatorPtr& alloc, OrtValue* p_mlvalue) {
static void CreateTensorMLValueOwned(const OrtPybindSingleUseAllocatorPtr& pybind_alloc, const AllocatorPtr& alloc, OrtValue* p_mlvalue) {
auto npy_type = PyArray_TYPE(pybind_alloc->GetContiguous());
TensorShape shape = GetArrayShape(pybind_alloc->GetContiguous());
auto element_type = NumpyToOnnxRuntimeTensorType(npy_type);
@ -389,9 +415,9 @@ std::string _get_type_name(std::string&) {
#if !defined(DISABLE_ML_OPS)
template <typename KeyType, typename ValueType, typename KeyGetterType, typename ValueGetterType>
void CreateMapMLValue_LoopIntoMap(Py_ssize_t& pos, PyObject*& key, const std::string& name_input, PyObject*& value,
PyObject* item, std::map<KeyType, ValueType>& current,
KeyGetterType keyGetter, ValueGetterType valueGetter) {
static void CreateMapMLValue_LoopIntoMap(Py_ssize_t& pos, PyObject*& key, const std::string& name_input, PyObject*& value,
PyObject* item, std::map<KeyType, ValueType>& current,
KeyGetterType keyGetter, ValueGetterType valueGetter) {
KeyType ckey;
ValueType cvalue;
do {
@ -427,9 +453,9 @@ void CreateMapMLValue_LoopIntoMap(Py_ssize_t& pos, PyObject*& key, const std::st
}
template <typename KeyType, typename ValueType, typename KeyGetterType, typename ValueGetterType>
void CreateMapMLValue_Map(Py_ssize_t& pos, PyObject*& key, const std::string& name_input, PyObject*& value,
PyObject* item, AllocatorPtr /*alloc*/, OrtValue* p_mlvalue, KeyGetterType keyGetter,
ValueGetterType valueGetter) {
static void CreateMapMLValue_Map(Py_ssize_t& pos, PyObject*& key, const std::string& name_input, PyObject*& value,
PyObject* item, AllocatorPtr /*alloc*/, OrtValue* p_mlvalue, KeyGetterType keyGetter,
ValueGetterType valueGetter) {
std::unique_ptr<std::map<KeyType, ValueType>> dst;
dst = onnxruntime::make_unique<std::map<KeyType, ValueType>>();
CreateMapMLValue_LoopIntoMap(pos, key, name_input, value, item, *dst, keyGetter, valueGetter);
@ -455,8 +481,8 @@ void CreateMapMLValue_VectorMap(Py_ssize_t& pos, PyObject*& key, const std::stri
DataTypeImpl::GetType<std::vector<std::map<KeyType, ValueType>>>()->GetDeleteFunc());
}
void CreateMapMLValue_AgnosticMap(Py_ssize_t& pos, PyObject*& key, const std::string& name_input, PyObject*& value,
PyObject* iterator, PyObject* item, AllocatorPtr alloc, OrtValue* p_mlvalue) {
static void CreateMapMLValue_AgnosticMap(Py_ssize_t& pos, PyObject*& key, const std::string& name_input, PyObject*& value,
PyObject* iterator, PyObject* item, AllocatorPtr alloc, OrtValue* p_mlvalue) {
// If iterator is NULL, it returns a single Map,
// if is not NULL, it returns a VectorMap.
auto int64Getter = [](PyObject* obj, int64_t& value) -> bool {
@ -526,8 +552,8 @@ void CreateMapMLValue_AgnosticMap(Py_ssize_t& pos, PyObject*& key, const std::st
}
}
void CreateMapMLValue_AgnosticVectorMap(PyObject* iterator, PyObject* item, AllocatorPtr alloc,
const std::string& name_input, OrtValue* p_mlvalue) {
static void CreateMapMLValue_AgnosticVectorMap(PyObject* iterator, PyObject* item, AllocatorPtr alloc,
const std::string& name_input, OrtValue* p_mlvalue) {
// CreateMapMLValue is called by CreateGenericTerableMLValue
// or CreateGenericMLValue which ensures
// item is a dictionary, no need to check type again.
@ -552,15 +578,15 @@ void CreateMapMLValue_AgnosticVectorMap(PyObject* iterator, PyObject* item, Allo
}
#endif
void CreateGenericIterableMLValue(PyObject* iterator, AllocatorPtr alloc, const std::string& name_input,
OrtValue* p_mlvalue) {
static void CreateGenericIterableMLValue(PyObject* iterator, AllocatorPtr alloc, const std::string& name_input,
OrtValue* p_mlvalue) {
PyObject* item;
OrtValue ml_value;
item = PyIter_Next(iterator);
if (item == NULL) {
throw std::runtime_error("Input '" + name_input + "' must not be empty.");
}
if (PyObjectCheck_Array(item)) {
if (PyObjectCheck_NumpyArray(item)) {
PyObject* pType = PyObject_Type(item);
PyObject* pStr = PyObject_Str(pType);
py::str spyType = py::reinterpret_borrow<py::str>(pStr);
@ -585,14 +611,19 @@ void CreateGenericIterableMLValue(PyObject* iterator, AllocatorPtr alloc, const
}
}
// Setting `use_numpy_data_memory` to `true` will ensure that the underlying numpy array buffer is directly used
// as the backing data buffer for the ORT Tensor where applicable (for numeric tensors)
// The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope
void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc, const std::string& name_input,
py::object& value, OrtValue* p_mlvalue) {
py::object& value, OrtValue* p_mlvalue, bool accept_only_numpy_array,
bool use_numpy_data_memory, MemCpyFunc mem_cpy_to_device) {
onnx::TypeProto type_proto;
if (PyObjectCheck_Array(value.ptr())) {
if (PyObjectCheck_NumpyArray(value.ptr())) {
// The most frequent case: input comes as an array.
PyArrayObject* arr = reinterpret_cast<PyArrayObject*>(value.ptr());
CreateTensorMLValue(alloc, name_input, arr, p_mlvalue);
} else if (PyList_Check(value.ptr()) &&
CreateTensorMLValue(alloc, name_input, arr, p_mlvalue, use_numpy_data_memory, mem_cpy_to_device);
} else if (!accept_only_numpy_array &&
PyList_Check(value.ptr()) &&
!CheckIfInputIsSequenceType(name_input, input_def_list, type_proto)) {
// This is not a sequence tensor. This is just a regular tensor fed through as a list.
ORT_ENFORCE(type_proto.tensor_type().has_elem_type(), "The graph is missing type information needed to construct the ORT tensor");
@ -612,12 +643,12 @@ void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const
// The allocator will own the array memory and will decrement the reference on Free()
// or when destroyed
auto pybind_allloc = std::make_shared<OrtPybindSingleUseAllocator>(arr, name_input, alloc->Info());
CreateTensorMLValueOwned(pybind_allloc, alloc, p_mlvalue);
} else if (PyList_Check(value.ptr())) {
auto pybind_alloc = std::make_shared<OrtPybindSingleUseAllocator>(arr, name_input, alloc->Info());
CreateTensorMLValueOwned(pybind_alloc, alloc, p_mlvalue);
} else if (!accept_only_numpy_array && PyList_Check(value.ptr())) {
auto* seq_tensors = reinterpret_cast<PyObject*>(value.ptr());
CreateSequenceOfTensors(alloc, name_input, input_def_list, seq_tensors, p_mlvalue);
} else if (PyDict_Check(value.ptr())) {
} else if (!accept_only_numpy_array && PyDict_Check(value.ptr())) {
#if !defined(DISABLE_ML_OPS)
CreateMapMLValue_AgnosticVectorMap((PyObject*)NULL, value.ptr(), alloc, name_input, p_mlvalue);
#else
@ -625,7 +656,13 @@ void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const
throw std::runtime_error("Map type is not supported in this build.");
#endif
} else {
} else if (!accept_only_numpy_array && strcmp(Py_TYPE(value.ptr())->tp_name, "OrtValue") == 0) {
// This is an OrtValue coming in directly from Python, so assign the underlying native OrtValue handle
// to the OrtValue object that we are going to use for Run().
// This should just increase the ref counts of the underlying shared_ptrs in the native OrtValue
// and the ref count will be decreased when the OrtValue used for Run() is destroyed upon exit.
*p_mlvalue = value.attr("_ortvalue").cast<OrtValue>();
} else if (!accept_only_numpy_array) {
auto iterator = PyObject_GetIter(value.ptr());
if (iterator == NULL) {
// The pype cannot be handled.
@ -646,6 +683,8 @@ void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const
throw;
}
Py_DECREF(iterator);
} else {
throw std::runtime_error("Unable to create OrtValue from the given python object");
}
}

View file

@ -21,14 +21,23 @@ namespace python {
namespace py = pybind11;
bool IsNumericNumpyType(int npy_type);
bool IsNumericNumpyArray(py::object& py_object);
int OnnxRuntimeTensorToNumpyType(const DataTypeImpl* tensor_type);
MLDataType NumpyTypeToOnnxRuntimeType(int numpy_type);
void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc, const std::string& name_input,
py::object& value, OrtValue* p_mlvalue);
using MemCpyFunc = void (*)(void*, const void*, size_t);
void CpuToCpuMemCpy(void*, const void*, size_t);
void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc,
const std::string& name_input, py::object& value, OrtValue* p_mlvalue,
bool accept_only_numpy_array = false, bool use_numpy_data_memory = true, MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy);
void GetPyObjFromTensor(const Tensor& rtensor, py::object& obj, const DataTransferManager* data_transfer_manager = nullptr);
void GetPyObjFromTensor(const Tensor& rtensor, py::object& obj,
const DataTransferManager* data_transfer_manager = nullptr,
const std::unordered_map<OrtDevice::DeviceType, MemCpyFunc>* mem_cpy_to_host_functions = nullptr);
template <class T>
struct DecRefFn {

View file

@ -139,6 +139,8 @@ struct OrtStatus {
#include "core/providers/cuda/cuda_provider_factory.h"
#include "core/providers/cuda/shared_inc/cuda_call.h"
#include "core/providers/cuda/cuda_execution_provider.h"
#include "core/providers/cuda/cuda_allocator.h"
OrtDevice::DeviceId cuda_device_id = 0;
OrtCudnnConvAlgoSearch cudnn_conv_algo_search = OrtCudnnConvAlgoSearch::EXHAUSTIVE;
size_t cuda_mem_limit = std::numeric_limits<size_t>::max();
@ -176,6 +178,12 @@ std::string nuphar_settings;
#define PYBIND_UNREFERENCED_PARAMETER(parameter) ((void)(parameter))
// Explicitly provide a definition for the static const var 'GPU' in the OrtDevice struct,
// GCC 4.x doesn't seem to define this and it breaks the pipelines based on CentOS as it uses
// GCC 4.x.
// (This static var is referenced in GetCudaToHostMemCpyFunction())
const OrtDevice::DeviceType OrtDevice::GPU;
namespace onnxruntime {
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_CPU(int use_arena);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_CUDA(OrtDevice::DeviceId device_id,
@ -186,9 +194,9 @@ std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_Tensor
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_MIGraphX(int device_id);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_Dnnl(int use_arena);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_NGraph(const char* ng_backend_type);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_OpenVINO(const char* device_type,
bool enable_vpu_fast_compile,
const char* device_id);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_OpenVINO(const char* device_type,
bool enable_vpu_fast_compile,
const char* device_id);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_Nuphar(bool, const char*);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_VITISAI(const char* backend_type, int device_id);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_ACL(int use_arena);
@ -265,11 +273,17 @@ void CustomOpLibrary::UnloadLibrary() {
#endif // !defined(ORT_MINIMAL_BUILD)
template <typename T>
void AddNonTensor(const OrtValue& val, std::vector<py::object>& pyobjs, const DataTransferManager* /*data_transfer_manager*/) {
static void AddNonTensor(const OrtValue& val, std::vector<py::object>& pyobjs,
const DataTransferManager* /*data_transfer_manager*/,
const std::unordered_map<OrtDevice::DeviceType, MemCpyFunc>* /*mem_cpy_to_host_functions*/) {
pyobjs.push_back(py::cast(val.Get<T>()));
}
void GetPyObjFromTensor(const Tensor& rtensor, py::object& obj, const DataTransferManager* data_transfer_manager) {
// In all cases, we may not have access to a DataTransferManager, hence the user may specify functions that
// pretty much does what a DataTransferManager does - copy data from device(s) to the host
void GetPyObjFromTensor(const Tensor& rtensor, py::object& obj,
const DataTransferManager* data_transfer_manager,
const std::unordered_map<OrtDevice::DeviceType, MemCpyFunc>* mem_cpy_to_host_functions) {
std::vector<npy_intp> npy_dims;
const TensorShape& shape = rtensor.Shape();
@ -282,25 +296,44 @@ void GetPyObjFromTensor(const Tensor& rtensor, py::object& obj, const DataTransf
obj = py::reinterpret_steal<py::object>(PyArray_SimpleNew(
shape.NumDimensions(), npy_dims.data(), numpy_type));
void* outPtr = static_cast<void*>(
void* out_ptr = static_cast<void*>(
PyArray_DATA(reinterpret_cast<PyArrayObject*>(obj.ptr())));
if (numpy_type != NPY_OBJECT) {
//if it is not cpu tensor, need to copy to host
if (rtensor.Location().device.Type() != OrtDevice::CPU) {
if (!data_transfer_manager)
throw std::runtime_error("GetPyObjFromTensor: data transfer manager is needed to convert non-CPU tensor to numpy array");
auto device_type = rtensor.Location().device.Type();
if (device_type != OrtDevice::CPU) {
if (!data_transfer_manager && !mem_cpy_to_host_functions)
throw std::runtime_error(
"GetPyObjFromTensor: Either data transfer manager or a "
"function to copy data to the host is needed to convert non-CPU tensor to numpy array");
static const OrtMemoryInfo cpu_alloc_info{onnxruntime::CPU, OrtDeviceAllocator};
std::vector<char> tensor_data_buffer{};
tensor_data_buffer.resize(rtensor.SizeInBytes());
ORT_THROW_IF_ERROR(CopyTensorDataToByteSpan(
*data_transfer_manager, rtensor, cpu_alloc_info, gsl::make_span(tensor_data_buffer)));
memcpy(outPtr, tensor_data_buffer.data(), dtype->Size() * shape.Size());
// Prefer DataTransferManager if available
if (data_transfer_manager) {
auto span = gsl::make_span<char>(reinterpret_cast<char*>(out_ptr), dtype->Size() * shape.Size());
ORT_THROW_IF_ERROR(CopyTensorDataToByteSpan(
*data_transfer_manager, rtensor, cpu_alloc_info, span));
} else {
auto mem_cpy_to_host = mem_cpy_to_host_functions->find(device_type);
ORT_ENFORCE(mem_cpy_to_host != mem_cpy_to_host_functions->end(),
"Unable to locate a function that can copy data to the host from the device");
ORT_ENFORCE(mem_cpy_to_host->second != 0,
"No function that can copy data to the host from the device provided");
mem_cpy_to_host->second(out_ptr, rtensor.DataRaw(), dtype->Size() * shape.Size());
}
} else
memcpy(outPtr, rtensor.DataRaw(dtype), dtype->Size() * shape.Size());
memcpy(out_ptr, rtensor.DataRaw(dtype), dtype->Size() * shape.Size());
} else {
// Handle string type.
py::object* outObj = static_cast<py::object*>(outPtr);
// Copying strings to cpu from device is currently not supported
ORT_ENFORCE(rtensor.Location().device.Type() == OrtDevice::CPU,
"Copying string tensors located on another device to the host is currently not supported");
py::object* outObj = static_cast<py::object*>(out_ptr);
const std::string* src = rtensor.template Data<std::string>();
for (int i = 0; i < rtensor.Shape().Size(); i++, src++) {
outObj[i] = py::cast(*src);
@ -322,49 +355,53 @@ static const char* GetDeviceName(const OrtDevice& device) {
}
template <>
void AddNonTensor<TensorSeq>(const OrtValue& val, std::vector<py::object>& pyobjs, const DataTransferManager* data_transfer_manager) {
void AddNonTensor<TensorSeq>(const OrtValue& val, std::vector<py::object>& pyobjs,
const DataTransferManager* data_transfer_manager,
const std::unordered_map<OrtDevice::DeviceType, MemCpyFunc>* mem_cpy_to_host_functions) {
const auto& seq_tensors = val.Get<TensorSeq>();
py::list py_list;
for (const auto& rtensor : seq_tensors) {
py::object obj;
GetPyObjFromTensor(rtensor, obj, data_transfer_manager);
GetPyObjFromTensor(rtensor, obj, data_transfer_manager, mem_cpy_to_host_functions);
py_list.append(obj);
}
pyobjs.push_back(py_list);
}
void AddNonTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs, const DataTransferManager* data_transfer_manager) {
static void AddNonTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs,
const DataTransferManager* data_transfer_manager,
const std::unordered_map<OrtDevice::DeviceType, MemCpyFunc>* mem_cpy_to_host_functions) {
// Should be in sync with core/framework/datatypes.h
auto val_type = val.Type();
if (val_type->IsTensorSequenceType()) {
AddNonTensor<TensorSeq>(val, pyobjs, data_transfer_manager);
AddNonTensor<TensorSeq>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else {
#if !defined(DISABLE_ML_OPS)
utils::ContainerChecker c_checker(val_type);
if (c_checker.IsMap()) {
if (c_checker.IsMapOf<std::string, std::string>()) {
AddNonTensor<MapStringToString>(val, pyobjs, data_transfer_manager);
AddNonTensor<MapStringToString>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else if (c_checker.IsMapOf<std::string, int64_t>()) {
AddNonTensor<MapStringToInt64>(val, pyobjs, data_transfer_manager);
AddNonTensor<MapStringToInt64>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else if (c_checker.IsMapOf<std::string, float>()) {
AddNonTensor<MapStringToFloat>(val, pyobjs, data_transfer_manager);
AddNonTensor<MapStringToFloat>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else if (c_checker.IsMapOf<std::string, double>()) {
AddNonTensor<MapStringToDouble>(val, pyobjs, data_transfer_manager);
AddNonTensor<MapStringToDouble>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else if (c_checker.IsMapOf<int64_t, std::string>()) {
AddNonTensor<MapInt64ToString>(val, pyobjs, data_transfer_manager);
AddNonTensor<MapInt64ToString>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else if (c_checker.IsMapOf<int64_t, int64_t>()) {
AddNonTensor<MapInt64ToInt64>(val, pyobjs, data_transfer_manager);
AddNonTensor<MapInt64ToInt64>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else if (c_checker.IsMapOf<int64_t, float>()) {
AddNonTensor<MapInt64ToFloat>(val, pyobjs, data_transfer_manager);
AddNonTensor<MapInt64ToFloat>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else if (c_checker.IsMapOf<int64_t, double>()) {
AddNonTensor<MapInt64ToDouble>(val, pyobjs, data_transfer_manager);
AddNonTensor<MapInt64ToDouble>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
}
} else {
if (c_checker.IsSequenceOf<std::map<std::string, float>>()) {
AddNonTensor<VectorMapStringToFloat>(val, pyobjs, data_transfer_manager);
AddNonTensor<VectorMapStringToFloat>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else if (c_checker.IsSequenceOf<std::map<int64_t, float>>()) {
AddNonTensor<VectorMapInt64ToFloat>(val, pyobjs, data_transfer_manager);
AddNonTensor<VectorMapInt64ToFloat>(val, pyobjs, data_transfer_manager, mem_cpy_to_host_functions);
} else {
throw std::runtime_error("Output is a non-tensor type which is not supported.");
}
@ -375,20 +412,22 @@ void AddNonTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs, c
}
}
void AddTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs, const DataTransferManager* data_transfer_manager) {
static void AddTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs,
const DataTransferManager* data_transfer_manager,
const std::unordered_map<OrtDevice::DeviceType, MemCpyFunc>* mem_cpy_to_host_functions) {
const Tensor& rtensor = val.Get<Tensor>();
py::object obj;
GetPyObjFromTensor(rtensor, obj, data_transfer_manager);
GetPyObjFromTensor(rtensor, obj, data_transfer_manager, mem_cpy_to_host_functions);
pyobjs.push_back(obj);
}
inline void RegisterExecutionProvider(InferenceSession* sess, onnxruntime::IExecutionProviderFactory& f) {
static inline void RegisterExecutionProvider(InferenceSession* sess, onnxruntime::IExecutionProviderFactory& f) {
auto p = f.CreateProvider();
OrtPybindThrowIfError(sess->RegisterExecutionProvider(std::move(p)));
}
// ordered by default priority from highest to lowest. kCpuExecutionProvider should always be last.
const std::vector<std::string>& GetAllProviders() {
static const std::vector<std::string>& GetAllProviders() {
static std::vector<std::string> all_providers = {kTensorrtExecutionProvider, kCudaExecutionProvider, kMIGraphXExecutionProvider,
kNGraphExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider,
kNupharExecutionProvider, kVitisAIExecutionProvider, kArmNNExecutionProvider,
@ -396,7 +435,7 @@ const std::vector<std::string>& GetAllProviders() {
return all_providers;
}
const std::vector<std::string>& GetAvailableProviders() {
static const std::vector<std::string>& GetAvailableProviders() {
auto InitializeProviders = []() {
std::vector<std::string> available_providers(std::begin(providers_available), std::end(providers_available));
return available_providers;
@ -413,7 +452,7 @@ const std::vector<std::string>& GetAvailableProviders() {
* (-1234, 43.21, +43.21 ... are not valid,
* 1234, +4321, 12 ... are valid)
*/
bool IsPositiveInteger(const std::string& s) {
static bool IsPositiveInteger(const std::string& s) {
if (s.length() == 0) {
return false;
}
@ -431,25 +470,25 @@ bool IsPositiveInteger(const std::string& s) {
return true;
}
bool IsCudaDeviceIdValid(InferenceSession* sess, int id) {
static bool IsCudaDeviceIdValid(const onnxruntime::logging::Logger& logger, int id) {
int num_devices = 0;
CUDA_CALL_THROW(cudaGetDeviceCount(&num_devices));
if (0 == num_devices) {
LOGS(*(sess->GetLogger()), WARNING) << "your system does not have a CUDA capable device.";
LOGS(logger, WARNING) << "your system does not have a CUDA capable device.";
return false;
}
if (id < 0 || id >= num_devices) {
LOGS(*(sess->GetLogger()), WARNING) << "cuda_device=" << id << " is invalid, must choose device ID between 0 and " << num_devices - 1;
LOGS(logger, WARNING) << "cuda_device=" << id << " is invalid, must choose device ID between 0 and " << num_devices - 1;
return false;
}
return true;
}
void UpdateCudaProviderOptions(InferenceSession* sess, onnxruntime::CUDAExecutionProviderInfo& options,
std::unordered_map<std::string, std::string> options_map) {
static void UpdateCudaProviderOptions(InferenceSession* sess, onnxruntime::CUDAExecutionProviderInfo& options,
std::unordered_map<std::string, std::string> options_map) {
std::unordered_map<std::string, std::string>::iterator it;
it = options_map.find("device_id");
@ -462,7 +501,7 @@ void UpdateCudaProviderOptions(InferenceSession* sess, onnxruntime::CUDAExecutio
throw std::runtime_error("Please provide device id with integer.");
}
if (!IsCudaDeviceIdValid(sess, device_id)) {
if (!IsCudaDeviceIdValid(*sess->GetLogger(), device_id)) {
throw std::runtime_error("Please provide available device id.");
}
options.device_id = device_id;
@ -511,6 +550,47 @@ void UpdateCudaProviderOptions(InferenceSession* sess, onnxruntime::CUDAExecutio
LOGS(*(sess->GetLogger()), INFO) << "cuda arean extend strategy is set to " << it->second;
}
}
static AllocatorPtr GetCudaAllocator(OrtDevice::DeviceId id) {
// Current approach is not thread-safe, but there are some bigger infra pieces to put together in order to make
// multi-threaded CUDA allocation work we need to maintain a per-thread CUDA allocator
static std::unordered_map<OrtDevice::DeviceId, AllocatorPtr> id_to_allocator_map;
if (id_to_allocator_map.find(id) == id_to_allocator_map.end()) {
// Use arena-based allocator
AllocatorCreationInfo default_memory_info(
[](OrtDevice::DeviceId id) {
return onnxruntime::make_unique<CUDAAllocator>(id, CUDA);
},
id,
true,
{cuda_mem_limit,
static_cast<int>(arena_extend_strategy),
-1, -1});
auto allocator = CreateAllocator(default_memory_info);
id_to_allocator_map.insert({id, allocator});
}
return id_to_allocator_map[id];
}
static void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes) {
CUDA_CALL_THROW(cudaMemcpy(dst, src, num_bytes, cudaMemcpyHostToDevice));
}
static void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) {
CUDA_CALL_THROW(cudaMemcpy(dst, src, num_bytes, cudaMemcpyDeviceToHost));
}
static const std::unordered_map<OrtDevice::DeviceType, MemCpyFunc>* GetCudaToHostMemCpyFunction() {
static std::unordered_map<OrtDevice::DeviceType, MemCpyFunc> map{
{OrtDevice::GPU, CudaToCpuMemCpy}};
return &map;
}
#endif
/*
@ -518,8 +598,8 @@ void UpdateCudaProviderOptions(InferenceSession* sess, onnxruntime::CUDAExecutio
*
* (note: currently only cuda EP supports this feature and rest of EPs use default options)
*/
void RegisterExecutionProviders(InferenceSession* sess, const std::vector<std::string>& provider_types,
ProviderOptionsMap& provider_options_map) {
static void RegisterExecutionProviders(InferenceSession* sess, const std::vector<std::string>& provider_types,
ProviderOptionsMap& provider_options_map) {
PYBIND_UNREFERENCED_PARAMETER(provider_options_map);
for (const std::string& type : provider_types) {
@ -569,11 +649,12 @@ void RegisterExecutionProviders(InferenceSession* sess, const std::vector<std::s
bool enable_vpu_fast_compile = false;
std::string openvino_device_id;
auto it = provider_options_map.find(type);
if(it != provider_options_map.end()) {
for(auto option : it->second) {
if(option.first == "device_type") openvino_device_type = option.second;
if (it != provider_options_map.end()) {
for (auto option : it->second) {
if (option.first == "device_type")
openvino_device_type = option.second;
else if (option.first == "enable_vpu_fast_compile") {
if(option.second == "True") {
if (option.second == "True") {
enable_vpu_fast_compile = true;
} else if (option.second == "False") {
enable_vpu_fast_compile = false;
@ -581,8 +662,8 @@ void RegisterExecutionProviders(InferenceSession* sess, const std::vector<std::s
ORT_THROW("Invalid value passed for enable_vpu_fast_compile: ", option.second);
}
}
else if (option.first == "device_id") openvino_device_id = option.second;
} else if (option.first == "device_id")
openvino_device_id = option.second;
else {
ORT_THROW("Invalid OpenVINO EP option: ", option.first);
}
@ -636,9 +717,9 @@ void RegisterExecutionProviders(InferenceSession* sess, const std::vector<std::s
* {'ep1' -> option1, 'ep2' -> option2 ...}
*
*/
void GenerateProviderOptionsMap(const std::vector<std::string>& providers,
const ProviderOptionsVector& provider_options_vector,
ProviderOptionsMap& provider_options_map) {
static void GenerateProviderOptionsMap(const std::vector<std::string>& providers,
const ProviderOptionsVector& provider_options_vector,
ProviderOptionsMap& provider_options_map) {
if (provider_options_vector.empty() || providers.empty()) {
return;
}
@ -655,7 +736,7 @@ void GenerateProviderOptionsMap(const std::vector<std::string>& providers,
}
#if !defined(ORT_MINIMAL_BUILD)
void RegisterCustomOpDomainsAndLibraries(PyInferenceSession* sess, const PySessionOptions& so) {
static void RegisterCustomOpDomainsAndLibraries(PyInferenceSession* sess, const PySessionOptions& so) {
if (!so.custom_op_domains_.empty()) {
// Register all custom op domains that will be needed for the session
std::vector<OrtCustomOpDomain*> custom_op_domains;
@ -685,6 +766,25 @@ void InitializeSession(InferenceSession* sess, const std::vector<std::string>& p
OrtPybindThrowIfError(sess->Initialize());
}
static bool CheckIfTensor(const std::vector<const NodeArg*>& def_list,
const std::string& name,
/*out*/ onnx::TypeProto& type_proto) {
auto ret_it = std::find_if(std::begin(def_list), std::end(def_list),
[&name](const NodeArg* node_arg) { return name == node_arg->Name(); });
if (ret_it == std::end(def_list)) {
throw std::runtime_error("Failed to find NodeArg with name: " + name + " in the def list");
}
const auto* temp = (*ret_it)->TypeAsProto();
if (!temp) {
throw std::runtime_error("Corresponding type_proto is null");
} else {
type_proto = *temp;
}
return type_proto.has_tensor_type();
}
void addGlobalMethods(py::module& m, const Environment& env) {
m.def("get_default_session_options", &GetDefaultCPUSessionOptions, "Return a default session_options instance.");
m.def("get_session_initializer", &SessionObjectInitializer::Get, "Return a default session object initializer.");
@ -727,7 +827,7 @@ void addGlobalMethods(py::module& m, const Environment& env) {
return ie_core.GetAvailableDevices();
},
"Lists all OpenVINO device ids available.");
/*
/*
* The following APIs to set config options are deprecated. Use Session.set_providers() instead.
*/
m.def(
@ -762,25 +862,25 @@ void addGlobalMethods(py::module& m, const Environment& env) {
onnxruntime::CreateExecutionProviderFactory_NGraph("CPU"),
#endif
#ifdef USE_OPENVINO
onnxruntime::CreateExecutionProviderFactory_OpenVINO(openvino_device_type, false, "");
onnxruntime::CreateExecutionProviderFactory_OpenVINO(openvino_device_type, false, ""),
#endif
#ifdef USE_TENSORRT
onnxruntime::CreateExecutionProviderFactory_Tensorrt(0),
#endif
#ifdef USE_MIGRAPHX
onnxruntime::CreateExecutionProviderFactory_MIGraphX(0)
onnxruntime::CreateExecutionProviderFactory_MIGraphX(0),
#endif
#ifdef USE_VITISAI
onnxruntime::CreateExecutionProviderFactory_VitisAI("DPU", 0),
onnxruntime::CreateExecutionProviderFactory_VitisAI("DPU", 0),
#endif
#ifdef USE_ACL
onnxruntime::CreateExecutionProviderFactory_ACL(0)
onnxruntime::CreateExecutionProviderFactory_ACL(0),
#endif
#ifdef USE_ARMNN
onnxruntime::CreateExecutionProviderFactory_ArmNN(0)
onnxruntime::CreateExecutionProviderFactory_ArmNN(0),
#endif
#ifdef USE_DML
onnxruntime::CreateExecutionProviderFactory_DML(0)
onnxruntime::CreateExecutionProviderFactory_DML(0)
#endif
};
@ -954,89 +1054,268 @@ void addObjectMethods(py::module& m, Environment& env) {
.def_static("cuda", []() { return OrtDevice::GPU; })
.def_static("default_memory", []() { return OrtDevice::MemType::DEFAULT; });
py::class_<SessionIOBinding> binding(m, "SessionIOBinding");
binding
py::class_<OrtValue> ortvalue_binding(m, "OrtValue");
ortvalue_binding
// Factory method to create an OrtValue (Tensor) from the given Numpy object
// The Tensor allocates and manages its own memory (on the specified device) and copies data from the Numpy data buffer
.def_static("ortvalue_from_numpy", [](py::object& array_on_cpu, OrtDevice& device) {
if (!IsNumericNumpyArray(array_on_cpu)) {
throw std::runtime_error("Creation of OrtValues is currently only supported from non-string numpy arrays");
}
auto ml_value = onnxruntime::make_unique<OrtValue>();
// The tensor's memory is allocated on the CPU
if (GetDeviceName(device) == CPU) {
// InputDeflist is null because OrtValue creation is not tied to a specific model
// Likewise, there is no need to specify the name (as the name was previously used to lookup the def list)
CreateGenericMLValue(nullptr, GetAllocator(), "", array_on_cpu, ml_value.get(), true);
} else if (GetDeviceName(device) == CUDA) {
// The tensor's memory is allocated on CUDA
#ifdef USE_CUDA
if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) {
throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine.");
}
// InputDeflist is null because OrtValue creation is not tied to a specific model
// Likewise, there is no need to specify the name (as the name was previously used to lookup the def list)
// TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in CUDA
CreateGenericMLValue(nullptr, GetCudaAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToCudaMemCpy);
#else
throw std::runtime_error(
"Can't allocate memory on the CUDA device using this package of OnnxRuntime. "
"Please use the CUDA package of OnnxRuntime to use this feature.");
#endif
} else {
throw std::runtime_error("Unsupported device: Cannot place the OrtValue on this device");
}
return ml_value;
})
// Factory method to create an OrtValue (Tensor) from the given shape and element type with memory on the specified device
// The memory is left uninitialized
.def_static("ortvalue_from_shape_and_type", [](std::vector<int64_t>& shape, py::object& element_type, OrtDevice& device) {
PyArray_Descr* dtype;
if (!PyArray_DescrConverter(element_type.ptr(), &dtype)) {
throw std::runtime_error("Not a valid numpy type");
}
int type_num = dtype->type_num;
Py_DECREF(dtype);
if (!IsNumericNumpyType(type_num)) {
throw std::runtime_error("Creation of OrtValues is currently only supported from non-string numpy arrays");
}
auto ml_value = onnxruntime::make_unique<OrtValue>();
std::unique_ptr<Tensor> tensor;
// The tensor's memory is allocated on the CPU
if (GetDeviceName(device) == CPU) {
tensor = onnxruntime::make_unique<Tensor>(NumpyTypeToOnnxRuntimeType(type_num), shape, GetAllocator());
} else if (GetDeviceName(device) == CUDA) {
// The tensor's memory is allocated on CUDA
#ifdef USE_CUDA
if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) {
throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine.");
}
tensor = onnxruntime::make_unique<Tensor>(NumpyTypeToOnnxRuntimeType(type_num), shape, GetCudaAllocator(device.Id()));
#else
throw std::runtime_error(
"Can't allocate memory on the CUDA device using this package of OnnxRuntime. "
"Please use the CUDA package of OnnxRuntime to use this feature.");
#endif
} else {
throw std::runtime_error("Unsupported device: Cannot place the OrtValue on this device");
}
auto ml_tensor = DataTypeImpl::GetType<Tensor>();
ml_value->Init(tensor.release(),
ml_tensor,
ml_tensor->GetDeleteFunc());
return ml_value;
})
.def("data_ptr", [](OrtValue* ml_value) -> int64_t {
// TODO: Assumes that the OrtValue is a Tensor, make this generic to handle non-Tensors
ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are currently supported");
auto* tensor = ml_value->GetMutable<Tensor>();
if (tensor->Shape().Size() == 0) {
return 0;
}
// Should cover x86 and x64 platforms
return reinterpret_cast<int64_t>(tensor->MutableDataRaw());
})
.def("device_name", [](OrtValue* ml_value) -> std::string {
// TODO: Assumes that the OrtValue is a Tensor, make this generic to handle non-Tensors
ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are currently supported");
return std::string(GetDeviceName(ml_value->Get<Tensor>().Location().device));
})
.def("shape", [](OrtValue* ml_value) -> py::list {
// TODO: Assumes that the OrtValue is a Tensor, make this generic to handle non-Tensors
ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are currently supported");
py::list shape_arr;
const auto& dims = ml_value->Get<Tensor>().Shape().GetDims();
for (auto dim : dims) {
// For sequence tensors - we would append a list of dims to the outermost list
// For now only tensors are supported in OrtValue
shape_arr.append(dim);
}
return shape_arr;
})
.def("data_type", [](OrtValue* ml_value) -> std::string {
// TODO: Assumes that the OrtValue is a Tensor, make this generic to handle non-Tensors
ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are currently supported");
return DataTypeImpl::ToString(ml_value->Get<Tensor>().DataType());
})
.def("is_tensor", [](OrtValue* ml_value) -> bool {
return ml_value->IsTensor();
})
.def("numpy", [](OrtValue* ml_value) -> py::object {
ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are convertible to Numpy objects");
py::object obj;
#ifdef USE_CUDA
GetPyObjFromTensor(ml_value->Get<Tensor>(), obj, nullptr, GetCudaToHostMemCpyFunction());
#else
GetPyObjFromTensor(ml_value->Get<Tensor>(), obj, nullptr, nullptr);
#endif
return obj;
});
py::class_<SessionIOBinding> session_io_binding(m, "SessionIOBinding");
session_io_binding
.def(py::init([](PyInferenceSession* sess) {
auto sess_io_binding = onnxruntime::make_unique<SessionIOBinding>(sess->GetSessionHandle());
return sess_io_binding;
}))
.def("bind_input", [](SessionIOBinding* io_binding, const std::string& name, py::object arr_on_cpu) -> void {
OrtValue mlvalue;
.def("bind_input", [](SessionIOBinding* io_binding, const std::string& name, py::object& arr_on_cpu) -> void {
InferenceSession* sess = io_binding->GetInferenceSession();
auto px = sess->GetModelInputs();
if (!px.first.IsOK() || !px.second) {
throw std::runtime_error("Either failed to get model inputs from the session object or the input def list was null");
}
//Check if input is sequence of tensors
onnx::TypeProto type_proto;
// For now, limit binding support to only non-string Tensors
// TODO: Support non-tensors
const auto& def_list = *px.second;
auto ret_it = std::find_if(std::begin(def_list), std::end(def_list),
[&name](const NodeArg* node_arg) { return name == node_arg->Name(); });
if (ret_it == std::end(def_list)) {
throw std::runtime_error("Failed to find input with name: " + name + " in the model input def list");
}
const auto* temp = (*ret_it)->TypeAsProto();
if (!temp) {
throw std::runtime_error("Corresponding type_proto is null");
} else {
type_proto = *temp;
}
if (type_proto.has_sequence_type()) {
throw std::runtime_error("Cannot bind input to sequence of tensors");
onnx::TypeProto type_proto;
if (!CheckIfTensor(def_list, name, type_proto)) {
throw std::runtime_error("Only binding Tensors is currently supported");
}
if (PyDict_Check(arr_on_cpu.ptr())) {
throw std::runtime_error("Cannot bind input to dictionary type of values");
ORT_ENFORCE(type_proto.tensor_type().has_elem_type());
if (type_proto.tensor_type().elem_type() == onnx::TensorProto::STRING) {
throw std::runtime_error("Only binding non-string Tensors is currently supported");
}
CreateGenericMLValue(px.second, GetAllocator(), name, arr_on_cpu, &mlvalue);
auto status = io_binding->Get()->BindInput(name, mlvalue);
if (!status.IsOK())
OrtValue ml_value;
// Set the parameter `accept_only_numpy_array` to `true` (we only support binding Tensors)
CreateGenericMLValue(px.second, GetAllocator(), name, arr_on_cpu, &ml_value, true);
auto status = io_binding->Get()->BindInput(name, ml_value);
if (!status.IsOK()) {
throw std::runtime_error("Error when bind input: " + status.ErrorMessage());
}
})
.def("bind_input", [](SessionIOBinding* io_binding, const std::string& name, const OrtDevice& device, py::object element_type, std::vector<int64_t> shape, int64_t data_ptr) -> void {
.def("bind_input", [](SessionIOBinding* io_binding, const std::string& name, const OrtDevice& device, py::object& element_type, std::vector<int64_t>& shape, int64_t data_ptr) -> void {
ORT_ENFORCE(data_ptr != 0, "Pointer to data memory is not valid");
PyArray_Descr* dtype;
if (!PyArray_DescrConverter(element_type.ptr(), &dtype))
if (!PyArray_DescrConverter(element_type.ptr(), &dtype)) {
throw std::runtime_error("Not a valid numpy type");
}
int type_num = dtype->type_num;
Py_DECREF(dtype);
OrtMemoryInfo info(GetDeviceName(device), OrtDeviceAllocator, device);
std::unique_ptr<Tensor> p_tensor = onnxruntime::make_unique<Tensor>(NumpyTypeToOnnxRuntimeType(type_num), shape, (void*)data_ptr, info);
OrtValue mlvalue;
std::unique_ptr<Tensor> p_tensor =
onnxruntime::make_unique<Tensor>(NumpyTypeToOnnxRuntimeType(type_num), shape, reinterpret_cast<void*>(data_ptr), info);
mlvalue.Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
auto status = io_binding->Get()->BindInput(name, mlvalue);
if (!status.IsOK())
throw std::runtime_error("Error when bind input: " + status.ErrorMessage());
OrtValue ml_value;
ml_value.Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
auto status = io_binding->Get()->BindInput(name, ml_value);
if (!status.IsOK()) {
throw std::runtime_error("Error when binding input: " + status.ErrorMessage());
}
})
.def("bind_output", [](SessionIOBinding* io_binding, const std::string& name, const OrtDevice& device, py::object element_type, std::vector<int64_t> shape, int64_t data_ptr) -> void {
.def("bind_ortvalue_input", [](SessionIOBinding* io_binding, const std::string& name, OrtValue& ml_value) -> void {
auto status = io_binding->Get()->BindInput(name, ml_value);
if (!status.IsOK()) {
throw std::runtime_error("Error when binding input: " + status.ErrorMessage());
}
})
.def("bind_output", [](SessionIOBinding* io_binding, const std::string& name, const OrtDevice& device, py::object& element_type, std::vector<int64_t>& shape, int64_t data_ptr) -> void {
ORT_ENFORCE(data_ptr != 0, "Pointer to data memory is not valid");
InferenceSession* sess = io_binding->GetInferenceSession();
auto px = sess->GetModelOutputs();
if (!px.first.IsOK() || !px.second) {
throw std::runtime_error("Either failed to get model inputs from the session object or the input def list was null");
}
// For now, limit binding support to only non-string Tensors
// TODO: Support non-tensors
const auto& def_list = *px.second;
onnx::TypeProto type_proto;
if (!CheckIfTensor(def_list, name, type_proto)) {
throw std::runtime_error("Only binding Tensors is currently supported");
}
ORT_ENFORCE(type_proto.tensor_type().has_elem_type());
if (type_proto.tensor_type().elem_type() == onnx::TensorProto::STRING) {
throw std::runtime_error("Only binding non-string Tensors is currently supported");
}
PyArray_Descr* dtype;
if (!PyArray_DescrConverter(element_type.ptr(), &dtype))
if (!PyArray_DescrConverter(element_type.ptr(), &dtype)) {
throw std::runtime_error("Not a valid numpy type");
}
int type_num = dtype->type_num;
Py_DECREF(dtype);
OrtMemoryInfo info(GetDeviceName(device), OrtDeviceAllocator, device);
std::unique_ptr<Tensor> p_tensor = onnxruntime::make_unique<Tensor>(NumpyTypeToOnnxRuntimeType(type_num), shape, (void*)data_ptr, info);
OrtValue mlvalue;
mlvalue.Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
std::unique_ptr<Tensor> p_tensor = onnxruntime::make_unique<Tensor>(NumpyTypeToOnnxRuntimeType(type_num), shape, reinterpret_cast<void*>(data_ptr), info);
auto status = io_binding->Get()->BindOutput(name, mlvalue);
if (!status.IsOK())
throw std::runtime_error("Error when bind output: " + status.ErrorMessage());
OrtValue ml_value;
ml_value.Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
auto status = io_binding->Get()->BindOutput(name, ml_value);
if (!status.IsOK()) {
throw std::runtime_error("Error when binding output: " + status.ErrorMessage());
}
})
.def("bind_output", [](SessionIOBinding* io_binding, const std::string& name, const OrtDevice& device) -> void {
auto status = io_binding->Get()->BindOutput(name, device);
if (!status.IsOK())
throw std::runtime_error("Error when bind output: " + status.ErrorMessage());
if (!status.IsOK()) {
throw std::runtime_error("Error when binding output: " + status.ErrorMessage());
}
})
.def("bind_ortvalue_output", [](SessionIOBinding* io_binding, const std::string& name, OrtValue& ml_value) -> void {
auto status = io_binding->Get()->BindOutput(name, ml_value);
if (!status.IsOK()) {
throw std::runtime_error("Error when binding output: " + status.ErrorMessage());
}
})
.def("clear_binding_inputs", [](SessionIOBinding* io_binding) -> void {
io_binding->Get()->ClearInputs();
@ -1044,15 +1323,18 @@ void addObjectMethods(py::module& m, Environment& env) {
.def("clear_binding_outputs", [](SessionIOBinding* io_binding) -> void {
io_binding->Get()->ClearOutputs();
})
.def("get_outputs", [](SessionIOBinding* io_binding) -> std::vector<OrtValue>& {
return io_binding->Get()->GetOutputs();
})
.def("copy_outputs_to_cpu", [](SessionIOBinding* io_binding) -> std::vector<py::object> {
const std::vector<OrtValue>& outputs = io_binding->Get()->GetOutputs();
std::vector<py::object> rfetch;
rfetch.reserve(outputs.size());
for (const auto& _ : outputs) {
if (_.IsTensor()) {
AddTensorAsPyObj(_, rfetch, &io_binding->GetInferenceSession()->GetDataTransferManager());
AddTensorAsPyObj(_, rfetch, &io_binding->GetInferenceSession()->GetDataTransferManager(), nullptr);
} else {
AddNonTensorAsPyObj(_, rfetch, &io_binding->GetInferenceSession()->GetDataTransferManager());
AddNonTensorAsPyObj(_, rfetch, &io_binding->GetInferenceSession()->GetDataTransferManager(), nullptr);
}
}
return rfetch;
@ -1098,7 +1380,7 @@ Applies to session load, initialization, etc. Default is 0.)pbdoc")
R"pbdoc(Sets the number of threads used to parallelize the execution of the graph (across nodes). Default is 0 to let onnxruntime choose.)pbdoc")
.def_readwrite("execution_mode", &PySessionOptions::execution_mode,
R"pbdoc(Sets the execution mode. Default is sequential.)pbdoc")
.def_readwrite("execution_order", &PySessionOptions::execution_order,
.def_readwrite("execution_order", &PySessionOptions::execution_order,
R"pbdoc(Sets the execution order. Default is basic topological order.)pbdoc")
.def_property(
"graph_optimization_level",
@ -1383,9 +1665,9 @@ including arg name, arg type (contains both type and shape).)pbdoc")
rfetch.reserve(fetches.size());
for (auto _ : fetches) {
if (_.IsTensor()) {
AddTensorAsPyObj(_, rfetch, nullptr);
AddTensorAsPyObj(_, rfetch, nullptr, nullptr);
} else {
AddNonTensorAsPyObj(_, rfetch, nullptr);
AddNonTensorAsPyObj(_, rfetch, nullptr, nullptr);
}
}
return rfetch;

View file

@ -712,6 +712,43 @@ class TestInferenceSession(unittest.TestCase):
so3.register_custom_ops_library(shared_library)
sess3 = onnxrt.InferenceSession(custom_op_model, so3)
def testOrtValue(self):
def test_session_with_ortvalue_input(ortvalue):
sess = onnxrt.InferenceSession(get_name("mul_1.onnx"))
res = sess.run(["Y"], {"X": ortvalue})
self.assertTrue(np.array_equal(res[0], numpy_arr_output))
numpy_arr_input = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32)
numpy_arr_output = np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32)
ortvalue1 = onnxrt.OrtValue.ortvalue_from_numpy(numpy_arr_input)
self.assertEqual(ortvalue1.device_name(), "cpu")
self.assertEqual(ortvalue1.shape(), [3, 2])
self.assertEqual(ortvalue1.data_type(), "tensor(float)")
self.assertEqual(ortvalue1.is_tensor(), True)
self.assertTrue(np.array_equal(ortvalue1.numpy(), numpy_arr_input))
# Pass in the constructed OrtValue to a session via Run() and check results
test_session_with_ortvalue_input(ortvalue1)
# The constructed OrtValue should still be valid after being used in a session
self.assertTrue(np.array_equal(ortvalue1.numpy(), numpy_arr_input))
if 'CUDAExecutionProvider' in onnxrt.get_available_providers():
ortvalue2 = onnxrt.OrtValue.ortvalue_from_numpy(numpy_arr_input, 'cuda', 0)
self.assertEqual(ortvalue2.device_name(), "cuda")
self.assertEqual(ortvalue2.shape(), [3, 2])
self.assertEqual(ortvalue2.data_type(), "tensor(float)")
self.assertEqual(ortvalue2.is_tensor(), True)
self.assertTrue(np.array_equal(ortvalue2.numpy(), numpy_arr_input))
# Pass in the constructed OrtValue to a session via Run() and check results
test_session_with_ortvalue_input(ortvalue2)
# The constructed OrtValue should still be valid after being used in a session
self.assertTrue(np.array_equal(ortvalue2.numpy(), numpy_arr_input))
if __name__ == '__main__':
unittest.main()

View file

@ -1,82 +1,42 @@
import numpy as np
import onnx
import onnxruntime
import torch
import torch.nn as nn
import unittest
from onnx import numpy_helper
class AtomicModel(nn.Module):
def __init__(self, kernel_size, in_channels, out_channels):
super(AtomicModel, self).__init__()
self.conv = nn.Conv2d(kernel_size=kernel_size,
in_channels=in_channels,
out_channels=out_channels)
def forward(self, x):
return self.conv(x)
from helper import get_name
class TestIOBinding(unittest.TestCase):
def create_model_and_input(self):
kernel_size = 5
channels = 16
device = torch.device('cuda')
batch_size = 1
sample_dim = 256
def create_ortvalue_input_on_gpu(self):
return onnxruntime.OrtValue.ortvalue_from_numpy(np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32), 'cuda', 0)
model = AtomicModel(kernel_size=kernel_size, in_channels=channels, out_channels=channels)
model.eval()
model.to(device)
# Run Pytorch
touch_input = torch.randn(batch_size, channels, sample_dim, sample_dim).to(device)
torch_output = model(touch_input)
torch_output_vals = torch_output.cpu().detach().numpy()
# Run ORT
input_names = [ "input" ]
output_names = [ "output" ]
torch.onnx.export(model, touch_input, "model.onnx",
input_names=input_names, output_names=output_names,
dynamic_axes={"input":{0:"batch_size"}, "output":{0:"batch_size"}})
return touch_input, torch_output, torch_output_vals
def create_ortvalue_alternate_input_on_gpu(self):
return onnxruntime.OrtValue.ortvalue_from_numpy(np.array([[2.0, 4.0], [6.0, 8.0], [10.0, 12.0]], dtype=np.float32), 'cuda', 0)
def test_bind_input_only(self):
torch_input, torch_output, torch_output_vals = self.create_model_and_input()
def create_uninitialized_ortvalue_input_on_gpu(self):
return onnxruntime.OrtValue.ortvalue_from_shape_and_type([3, 2], np.float32, 'cuda', 0)
session = onnxruntime.InferenceSession('model.onnx')
io_binding = session.io_binding()
# Bind input to CUDA
io_binding.bind_input('input', torch_input.device.type, 0, np.float32, list(torch_input.size()), torch_input.data_ptr())
def create_numpy_input(self):
return np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32)
# Bind output to CPU
io_binding.bind_output('output')
# Invoke Run
session.run_with_iobinding(io_binding)
# Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here)
ort_output = io_binding.copy_outputs_to_cpu()[0]
def create_expected_output(self):
return np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32)
# Validate results
self.assertTrue(np.array_equal(torch_output_vals, ort_output))
def create_expected_output_alternate(self):
return np.array([[2.0, 8.0], [18.0, 32.0], [50.0, 72.0]], dtype=np.float32)
def test_bind_input_to_cpu_arr(self):
torch_input, torch_output, torch_output_vals = self.create_model_and_input()
input = self.create_numpy_input()
session = onnxruntime.InferenceSession('model.onnx')
session = onnxruntime.InferenceSession(get_name("mul_1.onnx"))
io_binding = session.io_binding()
# Bind Numpy object (input) to CUDA
io_binding.bind_cpu_input('input', torch_input.cpu().detach().numpy())
# Bind Numpy object (input) that's on CPU to wherever the model needs it
io_binding.bind_cpu_input('X', self.create_numpy_input())
# Bind output to CPU
io_binding.bind_output('output')
io_binding.bind_output('Y')
# Invoke Run
session.run_with_iobinding(io_binding)
@ -85,63 +45,125 @@ class TestIOBinding(unittest.TestCase):
ort_output = io_binding.copy_outputs_to_cpu()[0]
# Validate results
self.assertTrue(np.array_equal(torch_output_vals, ort_output))
self.assertTrue(np.array_equal(self.create_expected_output(), ort_output))
def test_bind_input_only(self):
input = self.create_ortvalue_input_on_gpu()
session = onnxruntime.InferenceSession(get_name("mul_1.onnx"))
io_binding = session.io_binding()
# Bind input to CUDA
io_binding.bind_input('X', 'cuda', 0, np.float32, [3, 2], input.data_ptr())
# Bind output to CPU
io_binding.bind_output('Y')
# Invoke Run
session.run_with_iobinding(io_binding)
# Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here)
ort_output = io_binding.copy_outputs_to_cpu()[0]
# Validate results
self.assertTrue(np.array_equal(self.create_expected_output(), ort_output))
def test_bind_input_and_preallocated_output(self):
torch_input, torch_output, torch_output_vals = self.create_model_and_input()
input = self.create_ortvalue_input_on_gpu()
session = onnxruntime.InferenceSession('model.onnx')
session = onnxruntime.InferenceSession(get_name("mul_1.onnx"))
io_binding = session.io_binding()
# Bind input to CUDA
io_binding.bind_input('input', torch_input.device.type, 0, np.float32, list(torch_input.size()), torch_input.data_ptr())
io_binding.bind_input('X', 'cuda', 0, np.float32, [3, 2], input.data_ptr())
# Bind output to CUDA
ort_output = torch.empty(list(torch_output.size())).cuda()
io_binding.bind_output('output', torch_output.device.type, 0, np.float32, list(torch_output.size()), ort_output.data_ptr())
output = self.create_uninitialized_ortvalue_input_on_gpu()
io_binding.bind_output('Y', 'cuda', 0, np.float32, [3, 2], output.data_ptr())
# Invoke Run
session.run_with_iobinding(io_binding)
# Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here)
ort_output_vals = io_binding.copy_outputs_to_cpu()[0]
ort_output_vals = io_binding.copy_outputs_to_cpu()[0]
# Validate results
self.assertTrue(np.array_equal(torch_output_vals, ort_output_vals))
self.assertTrue(np.array_equal(self.create_expected_output(), ort_output_vals))
# Validate if ORT actually wrote to pre-allocated buffer by copying the Torch allocated buffer
# to the host and validating its contents
ort_output_vals = ort_output.cpu().detach().numpy()
ort_output_vals_in_cpu = output.numpy()
# Validate results
self.assertTrue(np.array_equal(torch_output_vals, ort_output_vals))
self.assertTrue(np.array_equal(self.create_expected_output(), ort_output_vals_in_cpu))
def test_bind_input_and_non_preallocated_output(self):
torch_input, torch_output, torch_output_vals = self.create_model_and_input()
session = onnxruntime.InferenceSession('model.onnx')
session = onnxruntime.InferenceSession(get_name("mul_1.onnx"))
io_binding = session.io_binding()
# Bind input to CUDA
io_binding.bind_input('input', torch_input.device.type, 0, np.float32, list(torch_input.size()), torch_input.data_ptr())
io_binding.bind_input('X', 'cuda', 0, np.float32, [3, 2], self.create_ortvalue_input_on_gpu().data_ptr())
# Bind output to CUDA
# DISCLAIMER: This is only useful for ORT benchmarking as there is really no way to access the
# ORT allocated device memory for this bound output at this point as ORT doesn't provide
# a handle over this memory yet and to access the contents of this memory, we would have to copy contents
# over to the host (CPU) using copy_outputs_to_cpu() as done in the next steps.
# In future, we will provide a handle over the ORT allocated device memory for the bound output.
# For now, this is a useful tool for benchmarking as Run() doesn't include the device-host copy latency
# by doing this.
io_binding.bind_output('output', 'cuda')
io_binding.bind_output('Y', 'cuda')
# Invoke Run
session.run_with_iobinding(io_binding)
# Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here)
ort_output = io_binding.copy_outputs_to_cpu()[0]
# Validate results
self.assertTrue(np.array_equal(torch_output_vals, ort_output))
# This call returns an OrtValue which has data allocated by ORT on CUDA
ort_outputs = io_binding.get_outputs()
self.assertEqual(len(ort_outputs), 1)
self.assertEqual(ort_outputs[0].device_name(), "cuda")
# Validate results (by copying results to CPU by creating a Numpy object)
self.assertTrue(np.array_equal(self.create_expected_output(), ort_outputs[0].numpy()))
# We should be able to repeat the above process as many times as we want - try once more
ort_outputs = io_binding.get_outputs()
self.assertEqual(len(ort_outputs), 1)
self.assertEqual(ort_outputs[0].device_name(), "cuda")
# Validate results (by copying results to CPU by creating a Numpy object)
self.assertTrue(np.array_equal(self.create_expected_output(), ort_outputs[0].numpy()))
# Change the bound input and validate the results in the same bound OrtValue
# Bind alternate input to CUDA
io_binding.bind_input('X', 'cuda', 0, np.float32, [3, 2], self.create_ortvalue_alternate_input_on_gpu().data_ptr())
# Invoke Run
session.run_with_iobinding(io_binding)
# This call returns an OrtValue which has data allocated by ORT on CUDA
ort_outputs = io_binding.get_outputs()
self.assertEqual(len(ort_outputs), 1)
self.assertEqual(ort_outputs[0].device_name(), "cuda")
# Validate results (by copying results to CPU by creating a Numpy object)
self.assertTrue(np.array_equal(self.create_expected_output_alternate(), ort_outputs[0].numpy()))
def test_bind_input_and_bind_output_with_ortvalues(self):
session = onnxruntime.InferenceSession(get_name("mul_1.onnx"))
io_binding = session.io_binding()
# Bind ortvalue as input
input_ortvalue = self.create_ortvalue_input_on_gpu()
io_binding.bind_ortvalue_input('X', input_ortvalue)
# Bind ortvalue as output
output_ortvalue = self.create_uninitialized_ortvalue_input_on_gpu()
io_binding.bind_ortvalue_output('Y', output_ortvalue)
# Invoke Run
session.run_with_iobinding(io_binding)
# Inspect contents of output_ortvalue and make sure that it has the right contents
self.assertTrue(np.array_equal(self.create_expected_output(), output_ortvalue.numpy()))
# Bind another ortvalue as input
input_ortvalue_2 = self.create_ortvalue_alternate_input_on_gpu()
io_binding.bind_ortvalue_input('X', input_ortvalue_2)
# Invoke Run
session.run_with_iobinding(io_binding)
# Inspect contents of output_ortvalue and make sure that it has the right contents
self.assertTrue(np.array_equal(self.create_expected_output_alternate(), output_ortvalue.numpy()))
if __name__ == '__main__':

View file

@ -7,7 +7,7 @@ import sys
import os
from onnxruntime.capi import _pybind_state as C
from onnxruntime.capi.session import Session, InferenceSession, IOBinding
from onnxruntime.capi.onnxruntime_inference_collection import Session, InferenceSession, IOBinding
class TrainingSession(InferenceSession):

View file

@ -1321,28 +1321,9 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs):
cwd=cwd, dll_path=dll_path)
# For CUDA enabled builds test IOBinding feature
# Limit testing to Windows non-ARM builds for now
iobinding_test = False
if args.use_cuda and not (args.arm or args.arm64):
if args.use_cuda:
# We need to have Torch installed to test the IOBinding feature
# which currently uses Torch's allocator to allocate GPU memory for testing
iobinding_test = True
# Try install Torch on Windows
if is_windows():
log.info("Attempting to install Torch to test ORT's IOBinding feature")
install_torch()
try:
import torch # noqa
except ImportError as error:
iobinding_test = False
log.exception(error)
log.warning(
"Torch is not installed. "
"The IOBinding tests will be skipped as it requires Torch.")
if iobinding_test:
log.info("Testing IOBinding feature")
run_subprocess([sys.executable, 'onnxruntime_test_python_iobinding.py'], cwd=cwd, dll_path=dll_path)