OlafenwaMoses/ImageAI

Any idea how to convert trained custom detection model (yolov3_last.pt) to onnx format?

Ebipatei opened this issue · 1 comments

Following the code below, it returns this error: AttributeError: 'collections.OrderedDict' object has no attribute 'training'

import torch
device = 'cuda:0' if torch.cuda.is_available else 'cpu'

load_saved_model = torch.load('labelled/models/yolov3_labelled_last.pt')

dummy_input = torch.FloatTensor(1,3,416,416)

dummy_input = dummy_input.to(device)

torch.onnx.export(load_saved_model, dummy_input, 'yolov3_labelled.onnx')

Any pointers will be appreciated.

@Ebipatei I will work on adding this functionality internally in ImageAI.