add ATen support for bicubic interpolation (#19380)

### Description
<!-- Describe your changes. -->

Add ATen fallback support for bicubic interpolation algorithm.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Required for facebook/dinov2 model architecture as part of ONNX Runtime
integration with AML Vision models.
This commit is contained in:
Prathik Rao 2024-02-05 13:11:37 -08:00 committed by GitHub
parent aaf32fb1b1
commit d120104dcd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 48 additions and 1 deletions

View file

@ -241,6 +241,7 @@ class SymbolicShapeInference:
"upsample_nearest1d": self._infer_aten_upsample,
"upsample_nearest2d": self._infer_aten_upsample,
"upsample_nearest3d": self._infer_aten_upsample,
"upsample_bicubic2d": self._infer_aten_upsample,
}
self.run_ = True
self.suggested_merge_ = {}

View file

@ -241,7 +241,7 @@ def native_group_norm_gradient():
# are available for all versions, though they are not that convienent to use.
def _upsample_gradient(backward_fn, dims):
scales = ["" for _ in range(dims)]
if "bilinear" in backward_fn:
if "bicubic" in backward_fn:
scales = ["I(2)", *scales]
return [
("Shape", ["I(0)"], ["Shape_X"]),
@ -271,3 +271,8 @@ def upsample_nearest2d_gradient():
@register_gradient("org.pytorch.aten", "ATen", "upsample_nearest3d", "vec")
def upsample_nearest3d_gradient():
return _upsample_gradient("upsample_nearest3d_backward", 3)
@register_gradient("org.pytorch.aten", "ATen", "upsample_bicubic2d", "vec")
def upsample_bicubic2d_gradient():
return _upsample_gradient("upsample_bicubic2d_backward", 2)

View file

@ -808,3 +808,16 @@ def upsample_nearest2d(g, input, output_size, scale_factors):
@register_symbolic("upsample_nearest3d")
def upsample_nearest3d(g, input, output_size, scale_factors):
return _upsample_nearest(g, input, output_size, scale_factors, "upsample_nearest3d")
@register_symbolic("upsample_bicubic2d")
def upsample_bicubic2d(g, input, output_size, align_corners, scale_factors):
return g.op(
"org.pytorch.aten::ATen",
input,
output_size,
align_corners,
scale_factors,
operator_s="upsample_bicubic2d",
overload_name_s="vec",
)

View file

@ -1805,6 +1805,34 @@ def test_resize_grad_correctness_bilinear_2d(interpolate_size_scale, align_corne
_test_helpers.assert_values_are_close(ort_input.grad, pt_input.grad)
def test_aten_upsample_bicubic():
class _NeuralNetUpsampleBicubic(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, input):
return torch.nn.functional.interpolate(input, size=(8, 12), mode="bicubic")
device = "cuda"
pt_model = _NeuralNetUpsampleBicubic().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)
pt_input = torch.randn([2, 4, 6, 8], 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):