Issues with exporting Pytorch model as ONNX

1

I am following the instructions from this tutorial on Pytorch for exporting to ONNX, so I can use Netron visualization.

Unfortunately, there seems to be a problem with onnxscript, even though I have already installed it through conda.

torch_model = torch.load(file_path)
Traceback (most recent call last):
  File "/home/user/export.py", line 22, in <module>
    main(FILE_PATH, shape)
  File "/home/user/export.py", line 10, in main
    onnx_program = torch.onnx.dynamo_export(torch_model, torch_input)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/test/lib/python3.12/site-packages/torch/onnx/_internal/exporter.py", line 1499, in dynamo_export
    resolved_export_options = ResolvedExportOptions(ExportOptions(), model=model)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/test/lib/python3.12/site-packages/torch/onnx/_internal/exporter.py", line 365, in __init__
    from torch.onnx._internal.fx import (  # TODO: Prevent circular dep
  File "/home/user/miniconda3/envs/test/lib/python3.12/site-packages/torch/onnx/_internal/fx/diagnostics.py", line 11, in <module>
    import onnxscript  # type: ignore[import]
    ^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'onnxscript'

Has anybody run into this issue?

2

The stacktrace is telling you that onnxscript is not found in your virtual environment, as stated in the tutorial, you should install:

pip install onnx
pip install onnxscript

If you have already installed both in the correct virtual environment it should work. Otherwise you may try this solution, it worked for me (adapt the code for your needs):

import torch
example_input = torch.randn(1, 3, args.img_height, args.img_width).to(device)

        torch.onnx.export(
            model,
            example_input,
            export_path,
            export_params=True,
            opset_version=args.onnx_opset_version,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size'},
                'output': {0: 'batch_size'}
            }
        )