From 975bb56e8c559c4db4b6d0059b7b8862150bfb79 Mon Sep 17 00:00:00 2001 From: Jameson Miller Date: Tue, 19 Jul 2022 14:37:03 -0400 Subject: [PATCH] Eager mode - argmax_out: set output tensor (#12233) This change updates the implementation or te argmax_out operator to 1) set the output tensor correctly and 2) remove the unnecessary use of a temporary tensor to store intermediate result of onnx ArgMax operation. Previously, the argmax_out operator did not correctly update the out tensor - it replaced the OrtValue instead of the memory backing the OrtValue . To properly update the output tensor, we need to calculate the expected shape of the out tensor. We add the helper function calculate_reduction_shape to calculate the shape of the reduced tensor from the input tensor, dimension to reduce, and option to keep the reduced dimension or not. This is based on the utility functions in aten/src/ATen/native/ReduceOpsUtils.h in the PyTorch repository, but is tailored to be a bit more specific to our current needs. Notes: We considered just directly leveraging PyTorch's utility functions (e.g. get_reduction_shape) to calculate the shape of the reduced tensor from aten/src/ATen/native/ReduceOpsUtils.h in the PyTorch repository, but including this header file resulted in warnings around unused functions that we need to handle. As we only need a limited functionality at the moment, we instead implemented our own utility function to calculate the reduction shape for our specific current needs. If we need a utility function to more generally calculate the reduction shape, we could consider switching to leveraging the utility methods in PyTorch. --- orttraining/orttraining/eager/ort_aten.cpp | 53 +++++++++++++++---- orttraining/orttraining/eager/test/ort_ops.py | 52 +++++++++++++++++- 2 files changed, 95 insertions(+), 10 deletions(-) diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp index 113205322c..f48b884520 100644 --- a/orttraining/orttraining/eager/ort_aten.cpp +++ b/orttraining/orttraining/eager/ort_aten.cpp @@ -365,6 +365,37 @@ OrtValue CastToType(onnxruntime::ORTInvoker& invoker, const OrtValue& input, at: return output[0]; } +/* + * Utility method to calculate the resulting shape of tensor after a reduction operation. + * + * @param dimToReduce The dimension to reduce. If null, then shape is 0 dimension. + * @param keepdim Whether to retain dim or not. Ignored if dimToReduce is null. + */ +inline at::DimVector calculate_reduction_shape( + const at::Tensor& self, + c10::optional dimToReduce, + bool keepdim) { + at::DimVector shape; + + // If we have dim value, then reduce that dimension. + // else, return empty shape (corresponding to 0-D tensor) + if (dimToReduce.has_value()) { + shape = at::DimVector(self.sizes()); + int64_t effectiveDimToReduce = *dimToReduce; + at::maybe_wrap_dims_n(&effectiveDimToReduce, 1, self.dim()); + + if (keepdim) { + shape[effectiveDimToReduce] = 1; + } else { + shape.erase(shape.begin() + effectiveDimToReduce); + } + } else { + shape = at::DimVector(); + } + + return shape; +} + /* * Utility function for resizing output tensor * Only resizes if: @@ -807,17 +838,27 @@ at::Tensor& out) { auto ort_input_self = create_ort_value(invoker, dim.has_value() ? self : self.reshape({-1})); - // Remove this hand signature once the generator can support this one line below. int64_t l_axis = dim.has_value() ? *dim : 0; + bool keepdim_effective_value = dim.has_value() ? keepdim : false; NodeAttributes attrs(2); attrs["axis"] = create_ort_attribute( - "axis", l_axis, at::ScalarType::Int); + "axis", l_axis, at::ScalarType::Int); attrs["keepdims"] = create_ort_attribute( - "keepdims", keepdim, at::ScalarType::Int); + "keepdims", keepdim_effective_value, at::ScalarType::Bool); std::vector ort_outputs_0_ArgMax(1); + // Calculate the size of the out tensor, based on self tensor, dimension input, and keepdim input + auto shape = calculate_reduction_shape(self, dim, keepdim); + + resize_output(invoker, + dynamic_cast(out.unsafeGetTensorImpl()), + at::IntArrayRef{shape}); + + auto ort_input_out = create_ort_value(invoker, out); + ort_outputs_0_ArgMax[0] = ort_input_out; + auto status = invoker.Invoke("ArgMax", { std::move(ort_input_self), }, ort_outputs_0_ArgMax, &attrs); @@ -826,12 +867,6 @@ at::Tensor& out) { throw std::runtime_error( "ORT return failure status:" + status.ErrorMessage()); - at::TensorOptions tensor_options = out.options(); - - // generator also needs to do this to handle the out param! - out = aten_tensor_from_ort( - std::move(ort_outputs_0_ArgMax[0]), - tensor_options); return out; } diff --git a/orttraining/orttraining/eager/test/ort_ops.py b/orttraining/orttraining/eager/test/ort_ops.py index 317c7415ee..a93c364335 100644 --- a/orttraining/orttraining/eager/test/ort_ops.py +++ b/orttraining/orttraining/eager/test/ort_ops.py @@ -258,13 +258,63 @@ class OrtOpTests(unittest.TestCase): def test_argmax(self): device = self.get_device() - cpu_tensor = torch.rand(3, 5) + cpu_tensor = torch.rand(3, 5, 7, 8) ort_tensor = cpu_tensor.to(device) + + # Scenario: basic (no dim parameters) + cpu_result = torch.argmax(cpu_tensor) + ort_result = torch.argmax(ort_tensor) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert cpu_result.dim() == ort_result.dim() + + # Scenario: specify dim parameter cpu_result = torch.argmax(cpu_tensor, dim=1) ort_result = torch.argmax(ort_tensor, dim=1) assert torch.allclose(cpu_result, ort_result.cpu()) assert cpu_result.dim() == ort_result.dim() + # Scenario: specify dim and keepdim parameters + cpu_result = torch.argmax(cpu_tensor, dim=1, keepdim=True) + ort_result = torch.argmax(ort_tensor, dim=1, keepdim=True) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert cpu_result.dim() == ort_result.dim() + + # Scenario: specify negative dim value + cpu_result = torch.argmax(cpu_tensor, dim=-1) + ort_result = torch.argmax(ort_tensor, dim=-1) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert cpu_result.dim() == ort_result.dim() + + # Scenario: basic out (no dim parameters) + cpu_out_tensor = torch.tensor([], dtype=torch.long) + ort_out_tensor = cpu_out_tensor.to(device) + cpu_result = torch.argmax(cpu_tensor, out=cpu_out_tensor) + ort_result = torch.argmax(ort_tensor, out=ort_out_tensor) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert cpu_result.dim() == ort_result.dim() + assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu()) + assert cpu_out_tensor.dim() == ort_out_tensor.dim() + + # Scenario: out with dim parameter + cpu_out_tensor = torch.tensor([], dtype=torch.long) + ort_out_tensor = cpu_out_tensor.to(device) + cpu_result = torch.argmax(cpu_tensor, dim=1, out=cpu_out_tensor) + ort_result = torch.argmax(ort_tensor, dim=1, out=ort_out_tensor) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert cpu_result.dim() == ort_result.dim() + assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu()) + assert cpu_out_tensor.dim() == ort_out_tensor.dim() + + # Scenario: out with dim and keepdim parameters + cpu_out_tensor = torch.tensor([], dtype=torch.long) + ort_out_tensor = cpu_out_tensor.to(device) + cpu_result = torch.argmax(cpu_tensor, dim=1, keepdim=True, out=cpu_out_tensor) + ort_result = torch.argmax(ort_tensor, dim=1, keepdim=True, out=ort_out_tensor) + assert torch.allclose(cpu_result, ort_result.cpu()) + assert cpu_result.dim() == ort_result.dim() + assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu()) + assert cpu_out_tensor.dim() == ort_out_tensor.dim() + def test_masked_select(self): device = self.get_device() cpu_tensor = torch.randn(3, 4)