Search code examples
machine-learningpytorchyoloonnx

convert pytorch model to ONNX

How to convert a pytorch model to ONNX? I am trying to use this method on Python 3.7:

import torch

model = torch.load("./yolov7x.pt")

#torch.onnx.export(model, "yolo_v7x.onnx")

Even with the commented last line in the 3 lines of code, loading this errors out:

Traceback (most recent call last):
  File "C:\Users\convert_onx.py", line 5, in <module>
    model = torch.load("./yolov7x.pt")
  File "C:\Users\Python37\lib\site-packages\torch\serialization.py", line 594, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "C:\Users\Python37\lib\site-packages\torch\serialization.py", line 853, in _load
    result = unpickler.load()
ModuleNotFoundError: No module named 'models'

This is the git repo I am working with the Yolo Model 7x: enter image description here

Ultimate use case is to use this model on Intel's Open VINO toolkit that requires PyTorch models to be converted to ONYX.

Solution

  • In the Yolo v7 repo: https://github.com/WongKinYiu/yolov7

    Just use the Google Collab notebooks provided in the Yolo v7 repo to convert the pytorch to different model types...easy as pie!