Convert PyTorch YOLOv5s model to ONNX or tensorflow format from google colab

1

Hi, i successfully trained a YOLOv5s model (.pt) in google colab. Could anyone please help me to find a way of how to convert PyTorch model weight format (.pt) to ONNX or tensorflow format using google colab?

Here is my colab code i try:

from torch.autograd import Variable

import torch.onnx

import torchvision

import torch

dummy_input = Variable(torch.randn(16, 3, 416, 416))

model = torch.load('/content/yolov5/runs/train/yolov5s_results2/weights/best.pt')

torch.onnx.export(model, dummy_input, "/content/yolov5/runs/train/yolov5s_results2/weights/pytorch_model.onnx")

But after run the code i’ve got some errors:

AttributeError                            Traceback (most recent call last)

<ipython-input-18-cef2b6adb415> in <module>()
      6 dummy_input = Variable(torch.randn(16, 3, 416, 416))
      7 model = torch.load('/content/yolov5/runs/train/yolov5s_results/weights/best.pt')
----> 8 torch.onnx.export(model,dummy_input,"/content/yolov5/runs/train/yolov5s_results/weights/pytorch_model.onnx")

4 frames

/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py in select_model_mode_for_export(model, mode)
     36 def select_model_mode_for_export(model, mode):
     37     if not isinstance(model, torch.jit.ScriptFunction):
---> 38         is_originally_training = model.training
     39 
     40         if mode is None:

AttributeError: 'dict' object has no attribute 'training'

What am i doing wrong?

2

torch.load() only loads the “state_dict”, this is NOT the actual model.

but i’m a bit confused here, there is some code like:

model = torch.load("weights/yolov5s.pt")['model']

on the other hand, the usual way to load a pytorch model is:

model = SomeNetwork() # first construct an empty nn from *code*
model.load_state_dict( torch.load("my.pth")) # fill in data