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

handle 8 bit uint dlpack tensor (#7069)

Browse source
This commit is contained in:
Vincent Wang 2021-03-20 08:00:49 +08:00 committed by GitHub
parent 8d5bfdeb47
commit cec919bae9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 72 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
onnxruntime/python/dlpack_convertor.cc
View file

@ -118,13 +118,13 @@ OrtDevice GetOrtDevice(const DLContext& ctx) {
}
}
MLDataType GetOrtValueDataType(const DLDataType& dtype) {
MLDataType GetOrtValueDataType(const DLDataType& dtype, bool is_bool_tensor) {
if (dtype.lanes != 1) ORT_THROW("ORT does not support lanes != 1");
switch (dtype.code) {
case DLDataTypeCode::kDLUInt:
switch (dtype.bits) {
case 8:
return DataTypeImpl::GetType<uint8_t>();
return is_bool_tensor ? DataTypeImpl::GetType<bool>() : DataTypeImpl::GetType<uint8_t>();
case 16:
return DataTypeImpl::GetType<uint16_t>();
case 32:
@ -213,11 +213,11 @@ DLManagedTensor* OrtValueToDlpack(const OrtValue& ort_value) {
return &(ort_dlmanaged_tensor->tensor);
}
OrtValue DlpackToOrtValue(const DLManagedTensor* dlpack) {
OrtValue DlpackToOrtValue(const DLManagedTensor* dlpack, bool is_bool_tensor) {
// ORT only supports contiguous tensor for now.
ORT_ENFORCE(IsContiguousTensor(dlpack->dl_tensor), "ORT only supports contiguous tensor for now.");
OrtDevice device = GetOrtDevice(dlpack->dl_tensor.ctx);
MLDataType data_type = GetOrtValueDataType(dlpack->dl_tensor.dtype);
MLDataType data_type = GetOrtValueDataType(dlpack->dl_tensor.dtype, is_bool_tensor);
std::function<void(void*)> deleter = [dlpack](void*) { dlpack->deleter(const_cast<DLManagedTensor*>(dlpack)); };
OrtMemoryInfo info(GetOrtDeviceName(device), OrtDeviceAllocator, device, device.Id());
std::unique_ptr<Tensor> p_tensor = onnxruntime::make_unique<Tensor>(

5
onnxruntime/python/dlpack_convertor.h
View file

@ -12,7 +12,10 @@ namespace onnxruntime {
namespace python {
DLManagedTensor* OrtValueToDlpack(const OrtValue& ort_value);
OrtValue DlpackToOrtValue(const DLManagedTensor* dlpack);
// DLPack uses same config for both bool and unit8. Parameter is_bool_tensor is to
// tell ORT the data type when creating OrtValue.
OrtValue DlpackToOrtValue(const DLManagedTensor* dlpack, bool is_bool_tensor = false);
} // namespace python
} // namespace onnxruntime

4
onnxruntime/python/onnxruntime_pybind_state.cc
View file

@ -1343,9 +1343,9 @@ void addObjectMethods(py::module& m, Environment& env) {
return py::reinterpret_steal<py::object>(
PyCapsule_New(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor));
})
.def_static("from_dlpack", [](py::object data) {
.def_static("from_dlpack", [](py::object data, bool is_bool_tensor = false) {
DLManagedTensor* dlmanaged_tensor = (DLManagedTensor*)PyCapsule_GetPointer(data.ptr(), "dltensor");
OrtValue ort_value = DlpackToOrtValue(dlmanaged_tensor);
OrtValue ort_value = DlpackToOrtValue(dlmanaged_tensor, is_bool_tensor);
// Make sure this capsule will never be used again.
PyCapsule_SetName(data.ptr(), "used_dltensor");
return ort_value;

37
orttraining/orttraining/python/training/ortmodule.py
View file

@ -30,12 +30,16 @@ T = TypeVar('T', bound='Module')
ONNX_OPSET_VERSION = 12
def _ortvalue_to_dlpack(ortvalue):
return ortvalue._ortvalue.to_dlpack()
def _ortvalue_to_torch_tensor(ortvalue):
# PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8,
# and convert the config to torch.uint8 tensor duing from_dlpack().
# So we need to convert the torch tensor to torch.bool type if OrtValue is bool tensor.
torch_tensor = from_dlpack(ortvalue._ortvalue.to_dlpack())
return torch_tensor.to(torch.bool) if ortvalue.data_type() == 'tensor(bool)' else torch_tensor
def _ortvalue_from_dlpack(dlpack_tensor):
return OrtValue(C.OrtValue.from_dlpack(dlpack_tensor))
def _ortvalue_from_torch_tensor(torch_tensor):
return OrtValue(C.OrtValue.from_dlpack(to_dlpack(torch_tensor), torch_tensor.dtype == torch.bool))
class Verbosity(IntEnum):
@ -50,7 +54,7 @@ def _create_iobinding(io_binding, inputs, model, device):
'''Creates IO binding for a `model` inputs and output'''
for idx, value_info in enumerate(model.graph.input):
io_binding.bind_ortvalue_input(
value_info.name, _ortvalue_from_dlpack(to_dlpack(inputs[idx])))
value_info.name, _ortvalue_from_torch_tensor(inputs[idx]))
for value_info in model.graph.output:
io_binding.bind_output(value_info.name, device.type,
@ -68,19 +72,6 @@ def _check_same_device(device, argument_str, *args):
f"{argument_str} found on device {arg_device}, but expected it to be on module device {device}.")
def _ort_output_to_torch_tensor(ort_output):
# TODO: PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8,
# and convert the config to torch.uint8 tensor duing from_dlpack(). So a boolean tensor
# from forward graph outputs will be converted to torch.uint8 tensor. When this tensor
# is feeded to backward graph as input, it will cause data type mismatch issue during
# inference session running. We cannot change the from_dlpack() in PyTorch side, so we
# have to handle this specially, which will introduce a cast here and there is data copied.
# Always cast from torch.uint8 to torch.bool is not logically right, we need to check the
# real data type of the inputs in the backeard graph, and perform the cast only necessary.
tensor = from_dlpack(_ortvalue_to_dlpack(ort_output))
return tensor.to(torch.bool) if tensor.dtype == torch.uint8 else tensor
def _load_torch_allocator_cpp_extension(verbosity):
torch_cuda_allocator_addresses_cpp_source = """
#include <torch/extension.h>
@ -177,7 +168,7 @@ class ORTModule(torch.nn.Module):
# Run and return module outputs.
forward_outputs, run_id = self._training_session.run_forward(
self._training_io_binding, self._run_options)
user_outputs = tuple(_ort_output_to_torch_tensor(
user_outputs = tuple(_ortvalue_to_torch_tensor(
forward_output) for forward_output in forward_outputs)
ctx.run_id = run_id
@ -219,8 +210,8 @@ class ORTModule(torch.nn.Module):
elif not grad_output.is_contiguous():
grad_output = grad_output.contiguous()
contiguous_grad_outputs.append(grad_output)
backward_grad_output_ortvalue = [_ortvalue_from_dlpack(
to_dlpack(grad_output)) for grad_output in contiguous_grad_outputs]
backward_grad_output_ortvalue = [_ortvalue_from_torch_tensor(
grad_output) for grad_output in contiguous_grad_outputs]
# Run and get results
run_id = ctx.run_id
@ -235,14 +226,14 @@ class ORTModule(torch.nn.Module):
for input_name in self._onnx_graphs_info.user_input_names:
try:
# Append to the results the backward output for each input that required grad
results.append(_ort_output_to_torch_tensor(
results.append(_ortvalue_to_torch_tensor(
backward_outputs[self._input_names_require_grad.index(input_name)]))
except ValueError:
# input_name is not found in the self._input_names_require_grad list
# Append None to results for each input that did not require grad
results.append(None)
# Append gradients of initializer to results
results += [_ort_output_to_torch_tensor(backward_output)
results += [_ortvalue_to_torch_tensor(backward_output)
for backward_output in backward_outputs[num_user_input_grads:]]
# The OrtValue has a shared_ptr to the data.
# At this point there are two shared_ptrs to the data, one through the

48
orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
View file

@ -1311,3 +1311,51 @@ def test_forward_returns_none_type_as_output():
assert output['out'] is not None
assert output['none_output'] is None
def test_bool_input_and_output():
class NeuralNetBoolInputOutput(torch.nn.Module):
def __init__(self, input_size, num_classes):
super(NeuralNetBoolInputOutput, self).__init__()
self.fc = torch.nn.Linear(input_size, num_classes)
self.relu = torch.nn.ReLU()
def forward(self, condition, x1, x2):
out1 = self.relu(self.fc(torch.where(condition, x1, x2)))
out2 = torch.tensor(out1).to(torch.bool)
return out1, out2
device = 'cuda'
N, D_in, D_out = 64, 784, 10
model = NeuralNetBoolInputOutput(D_in, D_out).to(device)
model = ORTModule(model)
condition = torch.randint(2, (N, D_in), dtype=torch.bool, device=device)
x1 = torch.randn(N, D_in, device=device)
x2 = torch.randn(N, D_in, device=device)
y1, y2 = model(condition, x1, x2)
assert y1 is not None
assert y2 is not None and y2.dtype == torch.bool
def test_uint8_input_and_output():
class NeuralNetUInt8InputOutput(torch.nn.Module):
def __init__(self, input_size, num_classes):
super(NeuralNetUInt8InputOutput, self).__init__()
self.fc = torch.nn.Linear(input_size, num_classes)
self.relu = torch.nn.ReLU()
def forward(self, mask, x1, x2):
out1 = self.relu(self.fc(torch.where(mask == 1, x1, x2)))
out2 = torch.tensor(out1).to(torch.uint8)
return out1, out2
device = 'cuda'
N, D_in, D_out = 64, 784, 10
model = NeuralNetUInt8InputOutput(D_in, D_out).to(device)
model = ORTModule(model)
condition = torch.randint(2, (N, D_in), dtype=torch.uint8, device=device)
x1 = torch.randn(N, D_in, device=device)
x2 = torch.randn(N, D_in, device=device)
y1, y2 = model(condition, x1, x2)
assert y1 is not None
assert y2 is not None and y2.dtype == torch.uint8