Conversion Error with Ultralytics YOLO Model using ai-edge-torch

lxzheng commented 2 months ago

Description of the bug:

I am encountering an InternalTorchDynamoError while converting a YOLO model from Ultralytics to TFLite using the ai-edge-torch library. The error message indicates an issue with unbacked symbols during the conversion process.

Actual vs expected behavior:

Actual: Using ai-edge-torch to convert the model results in the aforementioned error, preventing the generation of a TFLite model.
Expected: ai-edge-torch should be able to successfully perform the model conversion to the TFLite format, just like Ultralytics' built-in method does ( model.export(format="tflite")).

Any other information you'd like to share?

code to reproduce:

!pip install ultralytics
!pip install -r https://raw.githubusercontent.com/google-ai-edge/ai-edge-torch/main/requirements.txt
!pip install ai-edge-torch

import torch
import torchvision
import ai_edge_torch

from ultralytics import YOLO

model = YOLO('yolov8n.pt') 

model.export(format="tflite")
nhwc_model = ai_edge_torch.to_channel_last_io(model.model, args=[0])
edge_model = ai_edge_torch.convert(nhwc_model.eval(), (torch.randn(1, 640, 640, 3),))
edge_model.export("yolov8n.tflite")

error message:

InternalTorchDynamoError                  Traceback (most recent call last)
[<ipython-input-9-2fff8e335461>](https://localhost:8080/#) in <cell line: 13>()
     11 model.export(format="tflite")
     12 nhwc_model = ai_edge_torch.to_channel_last_io(model.model, args=[0])
---> 13 edge_model = ai_edge_torch.convert(nhwc_model.eval(), (torch.randn(1, 640, 640, 3),))
     14 edge_model.export("yolov8n.tflite")

122 frames
[/usr/local/lib/python3.10/dist-packages/torch/fx/experimental/symbolic_shapes.py](https://localhost:8080/#) in compute_unbacked_bindings(shape_env, example_value, old_example_value, peek)
    635                 else ""
    636             )
--> 637             raise PendingUnbackedSymbolNotFound(
    638                 f"Pending unbacked symbols {pending} not in returned outputs {example_value} {extra}.\n"
    639                 "Did you accidentally call new_dynamic_size() or item() more times "

InternalTorchDynamoError: Pending unbacked symbols {zuf0} not in returned outputs FakeTensor(..., size=(6400, 1)) ((1, 1), 0).
Did you accidentally call new_dynamic_size() or item() more times than you needed to in your fake implementation?
For more help, see https://docs.google.com/document/d/1RWrH-3wLEpzR9kCS6gGBNen_-Fs-8PVbWWFE5AcgeWE/edit

from user code:
   File "/usr/local/lib/python3.10/dist-packages/ai_edge_torch/convert/to_channel_last_io.py", line 51, in forward
    outputs = self.wrapped(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/nn/tasks.py", line 102, in forward
    return self.predict(x, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/nn/tasks.py", line 120, in predict
    return self._predict_once(x, profile, visualize, embed)
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/nn/tasks.py", line 141, in _predict_once
    x = m(x)  # run
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/nn/modules/head.py", line 61, in forward
    y = self._inference(x)
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/nn/modules/head.py", line 94, in _inference
    self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/utils/tal.py", line 307, in make_anchors
    stride_tensor.append(torch.full((h * w, 1), stride, dtype=dtype, device=device))

pkgoogle commented 2 months ago

Hi @lxzheng, can you try exporting via PyTorch first: https://pytorch.org/docs/stable/export.html, this lets us know if root cause is ai-edge-torch or PyTorch export. Thanks.

github-actions[bot] commented 2 months ago

Marking this issue as stale since it has been open for 7 days with no activity. This issue will be closed if no further activity occurs.

lxzheng commented 1 month ago

Hi @pkgoogle, Thank you for your suggestion. I attempted to export the model using the PyTorch torch.export API as recommended, but encountered the same issue.

pkgoogle commented 1 month ago

Hi @lxzheng, that model is not torch exportable... you can see that here in this gist. You can find more information here. Please raise an issue with PyTorch or use a different model, thanks.

github-actions[bot] commented 1 month ago

Marking this issue as stale since it has been open for 7 days with no activity. This issue will be closed if no further activity occurs.

github-actions[bot] commented 1 month ago

This issue was closed because it has been inactive for 14 days. Please post a new issue if you need further assistance. Thanks!

google-ai-edge / ai-edge-torch