aimet icon indicating copy to clipboard operation
aimet copied to clipboard

AIMET fails with YOLO models

Open Jamil opened this issue 1 year ago • 1 comments

When running

model = ultralytics.YOLO('yolov5nu.pt') # also with any ultralytics yolo model
dummy_input = torch.rand(1, 3, 640, 640).cuda()
ModelValidator.validate_model(model.model.cuda(), model_input=dummy_input)

I get:

2024-06-14 14:31:56,134 - Utils - INFO - Running validator check <function validate_for_reused_modules at 0x72bd5a507370>
2024-06-14 14:31:56,153 - Utils - ERROR - The following modules are used more than once in the model: ['model.0.act', 'model.9.m']
AIMET features are not designed to work with reused modules. Please redefine your model to use distinct modules for each instance.
2024-06-14 14:31:56,154 - Utils - INFO - Running validator check <function validate_for_missing_modules at 0x72bd5a507400>
2024-06-14 14:31:56,594 - Utils - ERROR - A connected graph failed to be built. This may prevent from AIMET features from being able to run on the model. Please address the errors shown.

And when preparing the model,

TraceError: Proxy object cannot be iterated. This can be attempted when the Proxy is used in a loop or as a *args or **kwargs function argument. See the torch.fx docs on pytorch.org for a more detailed explanation of what types of control flow can be traced, and check out the Proxy docstring for help troubleshooting Proxy iteration errors

Jamil avatar Jun 14 '24 21:06 Jamil

Hi @Jamil,

The model preparer issue is related to Torch FX tracing. To resolve this, you'll need to modify your code to ensure that the model is traceable by Torch FX.

Please refer to this doc for more details.

e-said avatar Jun 20 '24 15:06 e-said

Hi @Jamil as @e-said mentioned, this seems to be torch fx issue.

torch fx cannot handle branches that are present in Yolo model. In this case, you can skip model_preparer and directly go and create QuantizationSimMode

please refer to https://quic.github.io/aimet-pages/releases/latest/apiref/torch/quantsim.html#aimet_torch.QuantizationSimModel

Note that, you might have to decorate some of the modules as needed but will be a good start.

Also, you can try out aimet-onnx for this workflow where, you can export torch model to onnx and then use aimet_onnx.QuantizationSimModel to quantize this model. here's an example https://quic.github.io/aimet-pages/releases/latest/examples/onnx/quantization/quantsim.html

Closing this issue. but please feel free to open new issue depending on aimet-torch or aimet-onnx use directly

quic-bhushans avatar Apr 30 '25 23:04 quic-bhushans