I am trying to convert my yolov5 model .pt file to onnx format.It shows this error
Traceback (most recent call last): File "onnx.py", line 10, in <module> model = torch.load(weights) File "/home/nawinrajkumar/TrafficFlow/Traffic-flow-kgx/venv/lib/python3.8/site-packages/torch/serialization.py", line 789, in load return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args) File "/home/nawinrajkumar/TrafficFlow/Traffic-flow-kgx/venv/lib/python3.8/site-packages/torch/serialization.py", line 1131, in _load result = unpickler.load() File "/home/nawinrajkumar/TrafficFlow/Traffic-flow-kgx/venv/lib/python3.8/site-packages/torch/serialization.py", line 1124, in find_class return super().find_class(mod_name, name)
I used this code for the conversion of my model.
import torch
import torch.onnx
import sys
# Load your PyTorch model
weights = '/home/nawinrajkumar/TrafficFlow/Traffic-flow-kgx/yolov5s.pt'
model = torch.load(weights)
# Export the model to ONNX format
torch.onnx.export(model,
torch.randn(1, 3, 224, 224),
"model.onnx",
export_params=True)
What should I do to convert my .pt model to .onnx format. I also tried the export.py on the yolov5 repository and it doesn't help too.
