mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Support binding a graph output to a specific device via the Python binding (#4439)
This commit is contained in:
Hariharan Seshadri
GitHub
parent
aa06d308a6
commit
6d6b6b54a5
11 changed files with 194 additions and 67 deletions
|
|
@ -78,7 +78,16 @@ common::Status IOBinding::SynchronizeOutputs() {
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
common::Status IOBinding::BindOutput(const std::string& name, const OrtValue& ml_value, OrtDevice device) {
|
||||
common::Status IOBinding::BindOutput(const std::string& name, const OrtValue& ml_value) {
|
||||
// device value is ignored when ml_value is pre-allocated
|
||||
return BindOutputImpl(name, ml_value, {});
|
||||
}
|
||||
|
||||
common::Status IOBinding::BindOutput(const std::string& name, OrtDevice device) {
|
||||
return BindOutputImpl(name, {}, device);
|
||||
}
|
||||
|
||||
common::Status IOBinding::BindOutputImpl(const std::string& name, const OrtValue& ml_value, OrtDevice device) {
|
||||
auto rc = Contains(output_names_, name);
|
||||
if (rc.first) {
|
||||
outputs_[rc.second] = ml_value;
|
||||
|
|
|
|||
|
|
@ -60,13 +60,16 @@ class IOBinding {
|
|||
common::Status SynchronizeOutputs();
|
||||
|
||||
/**
|
||||
* Bind an output name to a provided OrtValue.
|
||||
* If the output is pre-allocated, the value in 'device' is not used.
|
||||
* If the output is not pre-allocated, information on the device it should be allocated on can be provided.
|
||||
*
|
||||
* @param device Device to allocate on if not pre-allocated. Default is CPU.
|
||||
* Bind an output name to a provided pre-allocated OrtValue.
|
||||
*/
|
||||
common::Status BindOutput(const std::string& name, const OrtValue& ml_value, OrtDevice device = {});
|
||||
common::Status BindOutput(const std::string& name, const OrtValue& ml_value);
|
||||
|
||||
/**
|
||||
* Bind an output name to a device.
|
||||
*
|
||||
* @param device Device to allocate the output on. Default is CPU.
|
||||
*/
|
||||
common::Status BindOutput(const std::string& name, OrtDevice device = {});
|
||||
|
||||
/**
|
||||
* This simply collects the outputs obtained after calling Run() inside the @param outputs.
|
||||
|
|
@ -74,9 +77,6 @@ class IOBinding {
|
|||
const std::vector<std::string>& GetOutputNames() const;
|
||||
std::vector<OrtValue>& GetOutputs();
|
||||
|
||||
// device info for outputs that are not pre-allocated
|
||||
const std::vector<OrtDevice>& GetOutputsDeviceInfo() const;
|
||||
|
||||
const std::vector<std::string>& GetInputNames() const;
|
||||
const std::vector<OrtValue>& GetInputs() const;
|
||||
|
||||
|
|
@ -105,5 +105,11 @@ class IOBinding {
|
|||
std::vector<OrtDevice> outputs_device_info_;
|
||||
|
||||
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(IOBinding);
|
||||
|
||||
// device info for all outputs. only used by InferenceSession if the output is not pre-allocated.
|
||||
const std::vector<OrtDevice>& GetOutputsDeviceInfo() const;
|
||||
|
||||
// The implementation for the BindOutput() overloads
|
||||
common::Status BindOutputImpl(const std::string& name, const OrtValue& ml_value, OrtDevice device);
|
||||
};
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -991,6 +991,10 @@ const SessionOptions& InferenceSession::GetSessionOptions() const {
|
|||
return session_options_;
|
||||
}
|
||||
|
||||
const DataTransferManager& InferenceSession::GetDataTransferManager() const {
|
||||
return data_transfer_mgr_;
|
||||
}
|
||||
|
||||
common::Status InferenceSession::CheckShapes(const std::string& input_name, const TensorShape& input_shape,
|
||||
const TensorShape& expected_shape) const {
|
||||
auto input_shape_sz = input_shape.NumDimensions();
|
||||
|
|
@ -1546,9 +1550,8 @@ common::Status InferenceSession::WaitForNotification(Notification* p_executor_do
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
SessionIOBinding::SessionIOBinding(InferenceSession* session) {
|
||||
SessionIOBinding::SessionIOBinding(InferenceSession* session) : sess_(session) {
|
||||
ORT_ENFORCE(session->NewIOBinding(&binding_).IsOK());
|
||||
sess_ = session;
|
||||
}
|
||||
|
||||
InferenceSession* SessionIOBinding::GetInferenceSession() {
|
||||
|
|
|
|||
|
|
@ -319,6 +319,11 @@ class InferenceSession {
|
|||
*/
|
||||
const SessionOptions& GetSessionOptions() const;
|
||||
|
||||
/*
|
||||
* Get the DataTransferManager associated with this session
|
||||
*/
|
||||
const DataTransferManager& GetDataTransferManager() const;
|
||||
|
||||
/**
|
||||
* Start profiling on this inference session. This simply turns on profiling events to be
|
||||
* recorded. A corresponding EndProfiling has to follow to write profiling data to a file.
|
||||
|
|
|
|||
|
|
@ -185,7 +185,7 @@ using namespace onnxruntime;
|
|||
using namespace onnxruntime::logging;
|
||||
|
||||
template <typename T>
|
||||
void AddNonTensor(const OrtValue& val, std::vector<py::object>& pyobjs) {
|
||||
void AddNonTensor(const OrtValue& val, std::vector<py::object>& pyobjs, const DataTransferManager* /*data_transfer_manager*/) {
|
||||
pyobjs.push_back(py::cast(val.Get<T>()));
|
||||
}
|
||||
|
||||
|
|
@ -242,47 +242,47 @@ static const char* GetDeviceName(const OrtDevice& device) {
|
|||
}
|
||||
|
||||
template <>
|
||||
void AddNonTensor<TensorSeq>(const OrtValue& val, std::vector<py::object>& pyobjs) {
|
||||
void AddNonTensor<TensorSeq>(const OrtValue& val, std::vector<py::object>& pyobjs, const DataTransferManager* data_transfer_manager) {
|
||||
const auto& seq_tensors = val.Get<TensorSeq>();
|
||||
py::list py_list;
|
||||
for (const auto& rtensor : seq_tensors) {
|
||||
py::object obj;
|
||||
GetPyObjFromTensor(rtensor, obj, nullptr);
|
||||
GetPyObjFromTensor(rtensor, obj, data_transfer_manager);
|
||||
py_list.append(obj);
|
||||
}
|
||||
pyobjs.push_back(py_list);
|
||||
}
|
||||
|
||||
void AddNonTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs) {
|
||||
void AddNonTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs, const DataTransferManager* data_transfer_manager) {
|
||||
// Should be in sync with core/framework/datatypes.h
|
||||
auto val_type = val.Type();
|
||||
if (val_type->IsTensorSequenceType()) {
|
||||
AddNonTensor<TensorSeq>(val, pyobjs);
|
||||
AddNonTensor<TensorSeq>(val, pyobjs, data_transfer_manager);
|
||||
} else {
|
||||
utils::ContainerChecker c_checker(val_type);
|
||||
if (c_checker.IsMap()) {
|
||||
if (c_checker.IsMapOf<std::string, std::string>()) {
|
||||
AddNonTensor<MapStringToString>(val, pyobjs);
|
||||
AddNonTensor<MapStringToString>(val, pyobjs, data_transfer_manager);
|
||||
} else if (c_checker.IsMapOf<std::string, int64_t>()) {
|
||||
AddNonTensor<MapStringToInt64>(val, pyobjs);
|
||||
AddNonTensor<MapStringToInt64>(val, pyobjs, data_transfer_manager);
|
||||
} else if (c_checker.IsMapOf<std::string, float>()) {
|
||||
AddNonTensor<MapStringToFloat>(val, pyobjs);
|
||||
AddNonTensor<MapStringToFloat>(val, pyobjs, data_transfer_manager);
|
||||
} else if (c_checker.IsMapOf<std::string, double>()) {
|
||||
AddNonTensor<MapStringToDouble>(val, pyobjs);
|
||||
AddNonTensor<MapStringToDouble>(val, pyobjs, data_transfer_manager);
|
||||
} else if (c_checker.IsMapOf<int64_t, std::string>()) {
|
||||
AddNonTensor<MapInt64ToString>(val, pyobjs);
|
||||
AddNonTensor<MapInt64ToString>(val, pyobjs, data_transfer_manager);
|
||||
} else if (c_checker.IsMapOf<int64_t, int64_t>()) {
|
||||
AddNonTensor<MapInt64ToInt64>(val, pyobjs);
|
||||
AddNonTensor<MapInt64ToInt64>(val, pyobjs, data_transfer_manager);
|
||||
} else if (c_checker.IsMapOf<int64_t, float>()) {
|
||||
AddNonTensor<MapInt64ToFloat>(val, pyobjs);
|
||||
AddNonTensor<MapInt64ToFloat>(val, pyobjs, data_transfer_manager);
|
||||
} else if (c_checker.IsMapOf<int64_t, double>()) {
|
||||
AddNonTensor<MapInt64ToDouble>(val, pyobjs);
|
||||
AddNonTensor<MapInt64ToDouble>(val, pyobjs, data_transfer_manager);
|
||||
}
|
||||
} else {
|
||||
if (c_checker.IsSequenceOf<std::map<std::string, float>>()) {
|
||||
AddNonTensor<VectorMapStringToFloat>(val, pyobjs);
|
||||
AddNonTensor<VectorMapStringToFloat>(val, pyobjs, data_transfer_manager);
|
||||
} else if (c_checker.IsSequenceOf<std::map<int64_t, float>>()) {
|
||||
AddNonTensor<VectorMapInt64ToFloat>(val, pyobjs);
|
||||
AddNonTensor<VectorMapInt64ToFloat>(val, pyobjs, data_transfer_manager);
|
||||
} else {
|
||||
throw std::runtime_error("Output is a non-tensor type which is not supported.");
|
||||
}
|
||||
|
|
@ -290,10 +290,10 @@ void AddNonTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs) {
|
|||
}
|
||||
}
|
||||
|
||||
void AddTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs) {
|
||||
void AddTensorAsPyObj(const OrtValue& val, std::vector<py::object>& pyobjs, const DataTransferManager* data_transfer_manager) {
|
||||
const Tensor& rtensor = val.Get<Tensor>();
|
||||
py::object obj;
|
||||
GetPyObjFromTensor(rtensor, obj);
|
||||
GetPyObjFromTensor(rtensor, obj, data_transfer_manager);
|
||||
pyobjs.push_back(obj);
|
||||
}
|
||||
|
||||
|
|
@ -460,13 +460,13 @@ void addGlobalMethods(py::module& m, const Environment& env) {
|
|||
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)
|
||||
#endif
|
||||
#ifdef USE_ARMNN
|
||||
onnxruntime::CreateExecutionProviderFactory_ArmNN(0)
|
||||
onnxruntime::CreateExecutionProviderFactory_ArmNN(0)
|
||||
#endif
|
||||
};
|
||||
|
||||
|
|
@ -642,7 +642,7 @@ void addObjectMethods(py::module& m, Environment& env) {
|
|||
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
|
||||
//Check if input is sequence of tensors
|
||||
onnx::TypeProto type_proto;
|
||||
const auto& def_list = *px.second;
|
||||
auto ret_it = std::find_if(std::begin(def_list), std::end(def_list),
|
||||
|
|
@ -654,7 +654,7 @@ void addObjectMethods(py::module& m, Environment& env) {
|
|||
if (!temp) {
|
||||
throw std::runtime_error("Corresponding type_proto is null");
|
||||
} else {
|
||||
type_proto = *temp;
|
||||
type_proto = *temp;
|
||||
}
|
||||
if (type_proto.has_sequence_type()) {
|
||||
throw std::runtime_error("Cannot bind input to sequence of tensors");
|
||||
|
|
@ -681,8 +681,8 @@ void addObjectMethods(py::module& m, Environment& env) {
|
|||
OrtValue mlvalue;
|
||||
|
||||
mlvalue.Init(p_tensor.release(),
|
||||
DataTypeImpl::GetType<Tensor>(),
|
||||
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
|
||||
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());
|
||||
|
|
@ -699,16 +699,15 @@ void addObjectMethods(py::module& m, Environment& env) {
|
|||
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());
|
||||
|
||||
DataTypeImpl::GetType<Tensor>(),
|
||||
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
|
||||
|
||||
auto status = io_binding->Get()->BindOutput(name, mlvalue);
|
||||
if (!status.IsOK())
|
||||
throw std::runtime_error("Error when bind output: " + status.ErrorMessage());
|
||||
})
|
||||
.def("bind_output", [](SessionIOBinding* io_binding, const std::string& name) -> void {
|
||||
OrtValue mlvalue;
|
||||
auto status = io_binding->Get()->BindOutput(name, mlvalue);
|
||||
.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());
|
||||
})
|
||||
|
|
@ -718,15 +717,15 @@ 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<py::object> {
|
||||
.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);
|
||||
AddTensorAsPyObj(_, rfetch, &io_binding->GetInferenceSession()->GetDataTransferManager());
|
||||
} else {
|
||||
AddNonTensorAsPyObj(_, rfetch);
|
||||
AddNonTensorAsPyObj(_, rfetch, &io_binding->GetInferenceSession()->GetDataTransferManager());
|
||||
}
|
||||
}
|
||||
return rfetch;
|
||||
|
|
@ -802,8 +801,7 @@ Applies to session load, initialization, etc. Default is 0.)pbdoc")
|
|||
},
|
||||
R"pbdoc(Graph optimization level for this session.)pbdoc")
|
||||
.def_readwrite("use_deterministic_compute", &SessionOptions::use_deterministic_compute,
|
||||
R"pbdoc(Whether to use deterministic compute. Default is false.)pbdoc")
|
||||
;
|
||||
R"pbdoc(Whether to use deterministic compute. Default is false.)pbdoc");
|
||||
|
||||
py::class_<RunOptions>(m, "RunOptions", R"pbdoc(Configuration information for a single Run.)pbdoc")
|
||||
.def(py::init())
|
||||
|
|
@ -955,9 +953,9 @@ including arg name, arg type (contains both type and shape).)pbdoc")
|
|||
rfetch.reserve(fetches.size());
|
||||
for (auto _ : fetches) {
|
||||
if (_.IsTensor()) {
|
||||
AddTensorAsPyObj(_, rfetch);
|
||||
AddTensorAsPyObj(_, rfetch, nullptr);
|
||||
} else {
|
||||
AddNonTensorAsPyObj(_, rfetch);
|
||||
AddNonTensorAsPyObj(_, rfetch, nullptr);
|
||||
}
|
||||
}
|
||||
return rfetch;
|
||||
|
|
|
|||
|
|
@ -228,17 +228,28 @@ class IOBinding:
|
|||
:param shape: output shape
|
||||
:param buffer_ptr: memory pointer to output data
|
||||
'''
|
||||
if device_type == 'cpu':
|
||||
self._iobinding.bind_output(name)
|
||||
|
||||
# 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 get_outputs(self):
|
||||
'''Obtain outputs.'''
|
||||
return self._iobinding.get_outputs()
|
||||
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()
|
||||
|
||||
def clear_binding_inputs(self):
|
||||
self._iobinding.clear_binding_inputs()
|
||||
|
|
|
|||
|
|
@ -308,7 +308,7 @@ void RunModelWithBindingMatMul(InferenceSession& session_object,
|
|||
|
||||
if (output_device) {
|
||||
// output should be allocated on specified device (if not preallocated here)
|
||||
io_binding->BindOutput("Y", output_ml_value, *output_device);
|
||||
io_binding->BindOutput("Y", *output_device);
|
||||
} else {
|
||||
io_binding->BindOutput("Y", output_ml_value);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -33,7 +33,8 @@ class TestIOBinding(unittest.TestCase):
|
|||
|
||||
# Run Pytorch
|
||||
touch_input = torch.randn(batch_size, channels, sample_dim, sample_dim).to(device)
|
||||
torch_output = model(touch_input).cpu().detach().numpy()
|
||||
torch_output = model(touch_input)
|
||||
torch_output_vals = torch_output.cpu().detach().numpy()
|
||||
|
||||
# Run ORT
|
||||
input_names = [ "input" ]
|
||||
|
|
@ -42,31 +43,106 @@ class TestIOBinding(unittest.TestCase):
|
|||
input_names=input_names, output_names=output_names,
|
||||
dynamic_axes={"input":{0:"batch_size"}, "output":{0:"batch_size"}})
|
||||
|
||||
return touch_input, torch_output
|
||||
return touch_input, torch_output, torch_output_vals
|
||||
|
||||
def test_bind_input_only(self):
|
||||
x, torch_output = self.create_model_and_input()
|
||||
torch_input, torch_output, torch_output_vals = self.create_model_and_input()
|
||||
|
||||
session = onnxruntime.InferenceSession('model.onnx')
|
||||
io_binding = session.io_binding()
|
||||
io_binding.bind_input('input', x.device.type, 0, np.float32, list(x.size()), x.data_ptr())
|
||||
|
||||
# Bind input to CUDA
|
||||
io_binding.bind_input('input', torch_input.device.type, 0, np.float32, list(torch_input.size()), torch_input.data_ptr())
|
||||
|
||||
# Bind output to CPU
|
||||
io_binding.bind_output('output')
|
||||
|
||||
# Invoke Run
|
||||
session.run_with_iobinding(io_binding)
|
||||
ort_output = io_binding.get_outputs()[0]
|
||||
|
||||
self.assertTrue(np.array_equal(torch_output, ort_output))
|
||||
|
||||
# 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))
|
||||
|
||||
def test_bind_input_to_cpu_arr(self):
|
||||
torch_input, torch_output = self.create_model_and_input()
|
||||
torch_input, torch_output, torch_output_vals = self.create_model_and_input()
|
||||
|
||||
session = onnxruntime.InferenceSession('model.onnx')
|
||||
io_binding = session.io_binding()
|
||||
|
||||
# Bind Numpy object (input) to CUDA
|
||||
io_binding.bind_cpu_input('input', torch_input.cpu().detach().numpy())
|
||||
|
||||
# Bind output to CPU
|
||||
io_binding.bind_output('output')
|
||||
|
||||
# Invoke Run
|
||||
session.run_with_iobinding(io_binding)
|
||||
ort_output = io_binding.get_outputs()[0]
|
||||
|
||||
self.assertTrue(np.array_equal(torch_output, ort_output))
|
||||
|
||||
# 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))
|
||||
|
||||
def test_bind_input_and_preallocated_output(self):
|
||||
torch_input, torch_output, torch_output_vals = self.create_model_and_input()
|
||||
|
||||
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())
|
||||
|
||||
# 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())
|
||||
|
||||
# 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]
|
||||
# Validate results
|
||||
self.assertTrue(np.array_equal(torch_output_vals, 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()
|
||||
# Validate results
|
||||
self.assertTrue(np.array_equal(torch_output_vals, ort_output_vals))
|
||||
|
||||
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')
|
||||
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())
|
||||
|
||||
# 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')
|
||||
|
||||
# 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))
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
|
|||
|
|
@ -102,7 +102,6 @@ def ort_training_session_run_helper(session, iobinding, inputs, input_descs, out
|
|||
torch_tensor = torch.zeros(output_desc.shape_, device=device,
|
||||
dtype=output_desc.eval_dtype_ if hasattr(output_desc, 'eval_dtype_')
|
||||
else output_desc.dtype_)
|
||||
|
||||
iobinding.bind_output(output_desc.name_, torch_tensor.device.type, get_device_index(device),
|
||||
dtype_torch_to_numpy(torch_tensor.dtype),
|
||||
list(torch_tensor.size()), torch_tensor.data_ptr())
|
||||
|
|
@ -195,6 +194,10 @@ def dtype_torch_to_numpy(torch_dtype):
|
|||
return np.int32
|
||||
elif torch_dtype == torch.int16 or torch_dtype == torch.short:
|
||||
return np.int16
|
||||
elif torch_dtype == torch.bool:
|
||||
return np.bool
|
||||
else:
|
||||
raise Exception("Torch type to numpy type mapping unavailable for: " + str(torch_dtype))
|
||||
|
||||
def wrap_for_input_match(model, loss_fn, input_names):
|
||||
import inspect
|
||||
|
|
|
|||
|
|
@ -1214,6 +1214,22 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs):
|
|||
[sys.executable, 'onnxruntime_test_python.py'],
|
||||
cwd=cwd, dll_path=dll_path)
|
||||
|
||||
iobinding_test = False
|
||||
# For CUDA enabled builds test IOBinding feature
|
||||
if args.use_cuda:
|
||||
iobinding_test = True
|
||||
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:
|
||||
run_subprocess([sys.executable, 'onnxruntime_test_python_iobinding.py'], cwd=cwd, dll_path=dll_path)
|
||||
|
||||
if not args.disable_ml_ops:
|
||||
run_subprocess([sys.executable, 'onnxruntime_test_python_mlops.py'], cwd=cwd, dll_path=dll_path)
|
||||
|
||||
|
|
|
|||
|
|
@ -113,8 +113,8 @@ if [ $DEVICE_TYPE = "Normal" ]; then
|
|||
${PYTHON_EXE} -m pip install sympy==1.1.1
|
||||
elif [ $DEVICE_TYPE = "gpu" ]; then
|
||||
${PYTHON_EXE} -m pip install sympy==1.1.1
|
||||
${PYTHON_EXE} -m pip install --upgrade --pre torch==1.6.0.dev20200610 torchvision==0.7.0.dev20200610 -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html
|
||||
if [[ $BUILD_EXTR_PAR = *--enable_training* ]]; then
|
||||
${PYTHON_EXE} -m pip install --upgrade --pre torch==1.6.0.dev20200610 torchvision==0.7.0.dev20200610 -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html
|
||||
${PYTHON_EXE} -m pip install transformers==v2.10.0
|
||||
# transformers requires sklearn
|
||||
${PYTHON_EXE} -m pip install sklearn
|
||||
|
|
|
|||
Loading…
Reference in a new issue