mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-28 20:11:22 +00:00
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:
parent
aaf32fb1b1
commit
d120104dcd
4 changed files with 48 additions and 1 deletions
|
|
@ -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_ = {}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in a new issue