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

[ORTModule] ATen Support for torch.nn.GroupNorm (#13293)

Browse source 
Model [huggingface's diffusers
library](https://github.com/huggingface/diffusers) has
torch.nn.GroupNorm which will be exported to sub-graph containing ONNX's
InstanceNormalization, which is lack of gradient. The implementation of
ORT's InstanceNormalization will call cuDNN's BatchNorm for part of
computation, which is not efficient compared to PyTorch's
implementation. This PR is to use ATen fallback to support this torch
module, including its forward and backward.
This commit is contained in:
Vincent Wang 2022-10-13 11:59:03 +08:00 committed by GitHub
parent 79ac0231a9
commit afb5f76770
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 118 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
orttraining/orttraining/core/graph/gradient_builder_base.cc
View file

@ -340,19 +340,46 @@ AttributeProto GradientBuilderBase::AttributeDefinitionToAttributeProto(
}
if (value.is_array()) {
ORT_ENFORCE(!attr_def.is_tensor, "1D tensor attribute is not supported for now.");
// Support only INTs for now.
switch (elem_type) {
case ONNX_NAMESPACE::TensorProto_DataType_INT64: {
std::vector<int64_t> int_array;
for (const auto& elem : value) {
ORT_ENFORCE(elem.is_number_integer());
int_array.emplace_back(static_cast<int64_t>(elem.get<int>()));
}
attr_proto = ONNX_NAMESPACE::MakeAttribute(name, int_array);
} break;
default:
ORT_ENFORCE(false, "List attribute of type ", std::to_string(elem_type), " is not supported for now.");
if (attr_def.is_tensor) {
// 1D tensor. Support only INT and BOOL types for now.
TensorProto tensor_proto;
switch (elem_type) {
case ONNX_NAMESPACE::TensorProto_DataType_INT64: {
std::vector<int64_t> int_array;
for (const auto& elem : value) {
ORT_ENFORCE(elem.is_number_integer());
int_array.emplace_back(static_cast<int64_t>(elem.get<int>()));
}
tensor_proto = ONNX_NAMESPACE::ToTensor<int64_t>(int_array);
tensor_proto.add_dims(static_cast<int64_t>(int_array.size()));
} break;
case ONNX_NAMESPACE::TensorProto_DataType_BOOL: {
std::vector<bool> bool_array;
for (const auto& elem : value) {
ORT_ENFORCE(elem.is_boolean());
bool_array.emplace_back(elem.get<bool>());
}
tensor_proto = ONNX_NAMESPACE::ToTensor<bool>(bool_array);
tensor_proto.add_dims(static_cast<int64_t>(bool_array.size()));
} break;
default:
ORT_ENFORCE(false, "Type ", std::to_string(elem_type), " is not supported for 1D tensor attribute.");
}
attr_proto = ONNX_NAMESPACE::MakeAttribute(name, tensor_proto);
} else {
// Support only INTs for now.
switch (elem_type) {
case ONNX_NAMESPACE::TensorProto_DataType_INT64: {
std::vector<int64_t> int_array;
for (const auto& elem : value) {
ORT_ENFORCE(elem.is_number_integer());
int_array.emplace_back(static_cast<int64_t>(elem.get<int>()));
}
attr_proto = ONNX_NAMESPACE::MakeAttribute(name, int_array);
} break;
default:
ORT_ENFORCE(false, "List attribute of type ", std::to_string(elem_type), " is not supported for now.");
}
}
} else if (attr_def.is_tensor) {
TensorProto tensor_proto;

13
orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py
View file

@ -222,3 +222,16 @@ def numpy_T_gradient():
{"operator": {"value": "numpy_T", "dtype": "string"}},
),
]
@register_gradient("org.pytorch.aten", "ATen", "native_group_norm", "")
def native_group_norm_gradient():
return [
("Constant", [], ["Const_0"], {"value": {"value": [True, True, True], "dtype": "bool", "is_tensor": True}}),
(
("ATen", "org.pytorch.aten"),
["GO(0)", "I(0)", "O(1)", "O(2)", "I(1)", "I(3)", "I(4)", "I(5)", "I(6)", "Const_0"],
["GI(0)", "GI(1)", "GI(2)"],
{"operator": {"value": "native_group_norm_backward", "dtype": "string"}},
),
]

30
orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py
View file

@ -663,3 +663,33 @@ def einsum_internal(g, equation, tensor_list):
# End of torch.einsum.
@register_symbolic("group_norm")
def group_norm(g, input, num_groups, weight, bias, eps, cudnn_enabled):
# Torch's group_norm's weight and bias are optional, its gradient has a bool[3] augment to indicate
# whether to compute the gradient for input, weight, bias. For simplicity of the gradient graph builder,
# we support only the case that weight and bias are not None.
from torch.onnx.symbolic_opset9 import group_norm as group_norm_generic
if weight is None or sym_help._is_none(weight) or bias is None or sym_help._is_none(bias):
return group_norm_generic(g, input, num_groups, weight, bias, eps, cudnn_enabled)
shape = g.op("Shape", input)
size = g.op("Size", input)
N = g.op("Gather", shape, g.op("Constant", value_t=torch.tensor(0, dtype=torch.long)), axis_i=0)
C = g.op("Gather", shape, g.op("Constant", value_t=torch.tensor(1, dtype=torch.long)), axis_i=0)
HxW = g.op("Div", size, g.op("Mul", N, C))
return g.op(
"org.pytorch.aten::ATen",
input,
weight,
bias,
N,
C,
HxW,
num_groups,
g.op("Cast", eps, to_i=1), # Python's float is float64.
operator_s="native_group_norm",
outputs=3,
)[0]

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

@ -1635,6 +1635,41 @@ def test_numpy_T(input_shape):
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
def test_aten_group_norm():
class NeuralNetGroupNorm(torch.nn.Module):
def __init__(self, num_groups, num_channels):
super(NeuralNetGroupNorm, self).__init__()
self.group_norm = torch.nn.GroupNorm(
num_groups=num_groups, num_channels=num_channels, eps=1e-5, affine=True
)
def forward(self, x, y):
return self.group_norm(x + y)
device = "cuda"
pt_model = NeuralNetGroupNorm(3, 6).to(device)
ort_model = ORTModule(copy.deepcopy(pt_model))
def run_step(model, x, y):
prediction = model(x, y)
prediction.sum().backward()
return prediction
# reset manual seed to reset the generator
torch.manual_seed(2333)
pt_x = torch.randn([20, 6, 10, 10], dtype=torch.float, device=device, requires_grad=True)
pt_y = torch.randn([20, 6, 10, 10], dtype=torch.float, device=device, requires_grad=True)
ort_x = copy.deepcopy(pt_x)
ort_y = copy.deepcopy(pt_y)
pt_prediction = run_step(pt_model, pt_x, pt_y)
ort_prediction = run_step(ort_model, ort_x, ort_y)
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
_test_helpers.assert_values_are_close(ort_x.grad, pt_x.grad)
_test_helpers.assert_values_are_close(ort_y.grad, pt_y.grad)
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model)
def test_gradient_correctness_cast_chain():
class NeuralNetCast(torch.nn.Module):
def __init__(self, D):