From 67150baa8d50e74f3a8a7b8e679a8d31eae4c0ed Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Thu, 20 Oct 2022 08:30:04 +0800 Subject: [PATCH] [ORTModule] ATen Support for aten::upsample_nearest (#13364) ATen support for aten::upsample_nearest, which is required for Huggingface's diffusers model training using ORTModule. --- .../python/tools/symbolic_shape_infer.py | 25 +++++++++- .../aten_op_executor/aten_op_executor.cc | 46 +++++++------------ .../ortmodule/_custom_gradient_registry.py | 30 ++++++++++++ .../ortmodule/_custom_op_symbolic_registry.py | 26 +++++++++++ .../python/orttraining_test_ortmodule_api.py | 35 ++++++++++++++ 5 files changed, 130 insertions(+), 32 deletions(-) diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 7921b602cd..ad35eff52a 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -212,6 +212,9 @@ class SymbolicShapeInference: "_adaptive_avg_pool2d": self._infer_aten_pool2d, "numpy_T": self._infer_Transpose, "native_group_norm": self._infer_aten_group_norm, + "upsample_nearest1d": self._infer_aten_upsample_nearest, + "upsample_nearest2d": self._infer_aten_upsample_nearest, + "upsample_nearest3d": self._infer_aten_upsample_nearest, } self.run_ = True self.suggested_merge_ = {} @@ -1366,12 +1369,30 @@ class SymbolicShapeInference: node.output[i], output_dtype, [ - N if N is not None else self._new_symbolic_dim_from_output(node, i, 0), - as_scalar(group) if group is not None else self._new_symbolic_dim_from_output(node, i, 1), + N if N is not None else str(self._new_symbolic_dim_from_output(node, i, 0)), + as_scalar(group) + if group is not None + else str(self._new_symbolic_dim_from_output(node, i, 1)), ], ) ) + def _infer_aten_upsample_nearest(self, node): + new_shape = None + input_shape = self._get_shape(node, 0) + if input_shape is not None: + new_shape = input_shape[:2] + output_size = self._try_get_value(node, 1) + if output_size is not None: + new_shape += [dim_size.item() for dim_size in output_size] + else: + rank = len(input_shape) + new_shape += [str(self._new_symbolic_dim_from_output(node, 0, i)) for i in range(2, rank)] + if node.output[0] and new_shape is not None: + output_dtype = self.known_vi_[node.input[0]].type.tensor_type.elem_type + vi = self.known_vi_[node.output[0]] + vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, new_shape)) + def _infer_BatchNormalization(self, node): self._propagate_shape_and_type(node) diff --git a/onnxruntime/python/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc b/onnxruntime/python/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc index e5b6dc8083..182f2368f5 100644 --- a/onnxruntime/python/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc +++ b/onnxruntime/python/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc @@ -6,42 +6,24 @@ #include #include -template +template c10::IValue ToIValue(const DLManagedTensor* dlpack, bool is_optional) { TORCH_INTERNAL_ASSERT((dlpack->dl_tensor.ndim == 0 && dlpack->dl_tensor.shape == nullptr) || (dlpack->dl_tensor.ndim == 1 && dlpack->dl_tensor.shape[0] == 1)); - T value = *reinterpret_cast(dlpack->dl_tensor.data); - return is_optional ? c10::IValue(c10::optional(value)) : c10::IValue(value); + TDst value = static_cast(*reinterpret_cast(dlpack->dl_tensor.data)); + return is_optional ? c10::IValue(c10::optional(value)) : c10::IValue(value); } -template +template c10::IValue ToListIValue(const DLManagedTensor* dlpack, bool is_optional) { TORCH_INTERNAL_ASSERT(dlpack->dl_tensor.ndim == 1); - const T* p_data = reinterpret_cast(dlpack->dl_tensor.data); - c10::List list_value; + const TSrc* p_data = reinterpret_cast(dlpack->dl_tensor.data); + c10::List list_value; size_t len = static_cast(dlpack->dl_tensor.shape[0]); for (size_t i = 0; i < len; i++) { - list_value.emplace_back(p_data[i]); + list_value.emplace_back(static_cast(p_data[i])); } - return is_optional ? c10::IValue(c10::optional>(list_value)) : c10::IValue(list_value); -} - -c10::IValue Int64ToBoolIValue(const DLManagedTensor* dlpack, bool is_list, bool is_optional) { - if (is_list) { - TORCH_INTERNAL_ASSERT(dlpack->dl_tensor.ndim == 1); - const int64_t* p_data = reinterpret_cast(dlpack->dl_tensor.data); - c10::List list_value; - size_t len = static_cast(dlpack->dl_tensor.shape[0]); - for (size_t i = 0; i < len; i++) { - list_value.emplace_back(static_cast(p_data[i])); - } - return is_optional ? c10::IValue(c10::optional>(list_value)) : c10::IValue(list_value); - } - - TORCH_INTERNAL_ASSERT((dlpack->dl_tensor.ndim == 0 && dlpack->dl_tensor.shape == nullptr) || - (dlpack->dl_tensor.ndim == 1 && dlpack->dl_tensor.shape[0] == 1)); - bool value = static_cast(*reinterpret_cast(dlpack->dl_tensor.data)); - return is_optional ? c10::IValue(c10::optional(value)) : c10::IValue(value); + return is_optional ? c10::IValue(c10::optional>(list_value)) : c10::IValue(list_value); } struct ATenOperator { @@ -78,22 +60,26 @@ struct ATenOperator { case c10::TypeKind::IntType: { TORCH_INTERNAL_ASSERT(dlpack->dl_tensor.dtype.code == DLDataTypeCode::kDLInt && dlpack->dl_tensor.dtype.bits == 64); - i_value = is_list ? ToListIValue(dlpack, is_optional) : ToIValue(dlpack, is_optional); + i_value = is_list ? ToListIValue(dlpack, is_optional) + : ToIValue(dlpack, is_optional); } break; case c10::TypeKind::FloatType: { TORCH_INTERNAL_ASSERT(dlpack->dl_tensor.dtype.code == DLDataTypeCode::kDLFloat && dlpack->dl_tensor.dtype.bits == 32); - i_value = is_list ? ToListIValue(dlpack, is_optional) : ToIValue(dlpack, is_optional); + // PyTorch's IValue doesn't support float, so we convert it to double. + i_value = + is_list ? ToListIValue(dlpack, is_optional) : ToIValue(dlpack, is_optional); } break; case c10::TypeKind::BoolType: { // In torch 1.8.1, exporter has bug which exports bool constant to int64 type tensor. // This bug has been fixed since torch 1.9.0. To make torch 1.8.1 work, add special handling here. if (dlpack->dl_tensor.dtype.code == DLDataTypeCode::kDLInt && dlpack->dl_tensor.dtype.bits == 64) { - i_value = Int64ToBoolIValue(dlpack, is_list, is_optional); + i_value = + is_list ? ToListIValue(dlpack, is_optional) : ToIValue(dlpack, is_optional); } else { TORCH_INTERNAL_ASSERT(dlpack->dl_tensor.dtype.code == DLDataTypeCode::kDLUInt && dlpack->dl_tensor.dtype.bits == 8); - i_value = is_list ? ToListIValue(dlpack, is_optional) : ToIValue(dlpack, is_optional); + i_value = is_list ? ToListIValue(dlpack, is_optional) : ToIValue(dlpack, is_optional); } } break; default: // TODO: will add more type support if needed. diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py index 1a1f850a1b..b7461781e7 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py @@ -235,3 +235,33 @@ def native_group_norm_gradient(): {"operator": {"value": "native_group_norm_backward", "dtype": "string"}}, ), ] + + +def _upsample_nearest_gradient(backward_fn): + return [ + ("Shape", ["I(0)"], ["Shape_X"]), + ( + ("ATen", "org.pytorch.aten"), + ["GO(0)", "I(1)", "Shape_X", "I(2)"], + ["GI(0)"], + { + "operator": {"value": backward_fn, "dtype": "string"}, + "overload_name": {"value": "vec", "dtype": "string"}, + }, + ), + ] + + +@register_gradient("org.pytorch.aten", "ATen", "upsample_nearest1d", "vec") +def upsample_nearest1d_gradient(): + return _upsample_nearest_gradient("upsample_nearest1d_backward") + + +@register_gradient("org.pytorch.aten", "ATen", "upsample_nearest2d", "vec") +def upsample_nearest2d_gradient(): + return _upsample_nearest_gradient("upsample_nearest2d_backward") + + +@register_gradient("org.pytorch.aten", "ATen", "upsample_nearest3d", "vec") +def upsample_nearest3d_gradient(): + return _upsample_nearest_gradient("upsample_nearest3d_backward") diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py index 3e7e93b73f..6c050a75b9 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py @@ -693,3 +693,29 @@ def group_norm(g, input, num_groups, weight, bias, eps, cudnn_enabled): operator_s="native_group_norm", outputs=3, )[0] + + +def _upsample_nearest(g, input, output_size, scale_factors, forward_fn): + return g.op( + "org.pytorch.aten::ATen", + input, + output_size, + scale_factors, + operator_s=forward_fn, + overload_name_s="vec", + ) + + +@register_symbolic("upsample_nearest1d") +def upsample_nearest1d(g, input, output_size, scale_factors): + return _upsample_nearest(g, input, output_size, scale_factors, "upsample_nearest1d") + + +@register_symbolic("upsample_nearest2d") +def upsample_nearest2d(g, input, output_size, scale_factors): + return _upsample_nearest(g, input, output_size, scale_factors, "upsample_nearest2d") + + +@register_symbolic("upsample_nearest3d") +def upsample_nearest3d(g, input, output_size, scale_factors): + return _upsample_nearest(g, input, output_size, scale_factors, "upsample_nearest3d") diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 8a1765ca61..85b25aae71 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -1691,6 +1691,41 @@ def test_aten_group_norm(): _test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model) +@pytest.mark.parametrize("input_rank", (3, 4, 5)) +@pytest.mark.parametrize("use_factor", (True, False)) +def test_aten_upsample_nearest(input_rank, use_factor): + class _NeuralNetUpsampleNearest(torch.nn.Module): + def __init__(self): + super(_NeuralNetUpsampleNearest, self).__init__() + + def forward(self, input): + return ( + torch.nn.functional.interpolate(input, scale_factor=2.0, mode="nearest") + if use_factor + else torch.nn.functional.interpolate(input, size=12, mode="nearest") + ) + + device = "cuda" + pt_model = _NeuralNetUpsampleNearest().to(device) + ort_model = ORTModule(copy.deepcopy(pt_model)) + + def run_step(model, input): + prediction = model(input) + prediction.sum().backward() + return prediction + + # reset manual seed to reset the generator + torch.manual_seed(2333) + input_size = [2 * (dim + 1) for dim in range(input_rank)] + pt_input = torch.randn(input_size, dtype=torch.float, device=device, requires_grad=True) + ort_input = copy.deepcopy(pt_input) + pt_prediction = run_step(pt_model, pt_input) + ort_prediction = run_step(ort_model, ort_input) + + _test_helpers.assert_values_are_close(ort_prediction, pt_prediction) + _test_helpers.assert_values_are_close(ort_input.grad, pt_input.grad) + + def test_gradient_correctness_cast_chain(): class NeuralNetCast(torch.nn.Module): def __init__(self, D):