Skip to content

[OnnxToTorch] Fix Resize op when ONNX exports dynamic spatial dims as 0 #4294

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

HalfBloodPrince010
Copy link

ONNX encodes dynamic spatial dimensions in Resize as 0. This pass just passes the spatial dimensions as is for next steps. For interpolate/resize op, [0,0] spatial dimension is not valid proposed size. This change replaces such 0 values with the corresponding runtime dimension from the input tensor, ensuring correct shape propagation in Torch-MLIR and preventing invalid 0-sized dimensions.

This fixes the issue: iree-org/iree#19501

Example for ONNX Model Export

>>> import torch
>>> import torch.nn as nn
>>> import torch.onnx
>>>
>>> class InterpolateModel(nn.Module):
...     def forward(self, x):
...         return nn.functional.interpolate(x, scale_factor=0.5, mode='bicubic', align_corners=False)
...
>>> model = InterpolateModel()
>>> dummy_input = torch.randn(11, 3, 32, 54)
>>>
>>> torch.onnx.export(
...     model,
...     dummy_input,
...     "interpolate_dynamic_cubic.onnx",
...     opset_version=11,
...     input_names=['pixel_values'],
...     output_names=['output'],
...     dynamic_axes={
...         'pixel_values': {0: 'batch', 2: 'height', 3: 'width'},
...         'output': {0: 'batch', 2: 'height', 3: 'width'},
...     }
... )

After exporting the model
>>>
>>> print("Exported interpolate_dynamic_cubic.onnx")
Exported interpolate_dynamic_cubic.onnx
>>> dynamic_input_model = onnx.load("interpolate_dynamic_cubic.onnx")
>>> for input_tensor in dynamic_input_model.graph.input:
...     shape = [dim.dim_value for dim in input_tensor.type.tensor_type.shape.dim]
...     print(f"{input_tensor.name}: {shape}")
...
pixel_values: [0, 3, 0, 0]

In IREE, one of the test using Hugging Model is failing because the interpolate/resize
dimensions are [0, 0]

  1. Input (pixel_values)
  2. Get Shape
  3. Gather indices 2 and 3 (which is spatial dimension height and width)
  4. Divide both by 16 (patch size)
  5. Concat them.

But input pixrl_values are all [0, 0 , 0, 0]

import onnx
model = onnx.load("model.onnx")

>>> for input_tensor in model.graph.input:
...     shape = [dim.dim_value for dim in input_tensor.type.tensor_type.shape.dim]
...     print(f"{input_tensor.name}: {shape}")
...
pixel_values: [0, 0, 0, 0]

ONNX encodes dynamic spatial dimensions in Resize as 0. This pass
just passes the spatial dimensions as it. For interpolate/resize
op, [0,0] spatial dimension is not valid proposed size.
This change replaces such 0 values with the corresponding
runtime dimension from the input tensor, ensuring correct shape
propagation in Torch-MLIR and preventing invalid 0-sized dimensions.

Signed-off-by: Prashanth Pujar <prashanth.pujar.usc@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant