convert pytorch model to ONNX-CodePudding

How to convert a pytorch model to ONNX? I am trying to use

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

CodePudding user response：

When you are loading the pickled model the source tree must match the one that used when the model was saved. So

ModuleNotFoundError: No module named 'models'

expects this directory to be in your python path:

Usually it is better to save weights as state_dict and keep the source code that can reconstruct the torch.nn.Module so then you can safely use:

model.load_state_dict(torch.load('weights.pt'))

CodePudding user response：

For what its worth I could only get this to work jerry rigging some code from another repo, plus the answer from @u1234x1234

import torch
from models.experimental import attempt_load
import os

weights ='yolov7x.pt'
device = torch.device('cpu')

# Load model
model = attempt_load(weights, map_location=device) 

torch.onnx.export(model, torch.zeros((1, 3, 640, 640)), 'yolov7.onnx', opset_version=12)

Appeared to work with this function attempt_load a file from this other pytorch user repo. I am unfamiliar with pytorch and dont know if this is just a common pytorch file or not...