From cec919bae991a5ba826e1f7282ce51fa30cbfa66 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Sat, 20 Mar 2021 08:00:49 +0800 Subject: [PATCH] handle 8 bit uint dlpack tensor (#7069) --- onnxruntime/python/dlpack_convertor.cc | 8 ++-- onnxruntime/python/dlpack_convertor.h | 5 +- .../python/onnxruntime_pybind_state.cc | 4 +- .../orttraining/python/training/ortmodule.py | 37 ++++++-------- .../python/orttraining_test_ortmodule_api.py | 48 +++++++++++++++++++ 5 files changed, 72 insertions(+), 30 deletions(-) diff --git a/onnxruntime/python/dlpack_convertor.cc b/onnxruntime/python/dlpack_convertor.cc index a4d11eec1c..eb92960494 100644 --- a/onnxruntime/python/dlpack_convertor.cc +++ b/onnxruntime/python/dlpack_convertor.cc @@ -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(); + return is_bool_tensor ? DataTypeImpl::GetType() : DataTypeImpl::GetType(); case 16: return DataTypeImpl::GetType(); 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 deleter = [dlpack](void*) { dlpack->deleter(const_cast(dlpack)); }; OrtMemoryInfo info(GetOrtDeviceName(device), OrtDeviceAllocator, device, device.Id()); std::unique_ptr p_tensor = onnxruntime::make_unique( diff --git a/onnxruntime/python/dlpack_convertor.h b/onnxruntime/python/dlpack_convertor.h index 659b6c91e3..98f1c8663c 100644 --- a/onnxruntime/python/dlpack_convertor.h +++ b/onnxruntime/python/dlpack_convertor.h @@ -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 diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 91507fa19f..5a26cfd4de 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -1343,9 +1343,9 @@ void addObjectMethods(py::module& m, Environment& env) { return py::reinterpret_steal( 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; diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index c1493b35fc..dbe3b2ae1c 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -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 @@ -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 diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 61e92d4a1a..0baa21c224 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -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