OlafenwaMoses / ImageAI

A python library built to empower developers to build applications and systems with self-contained Computer Vision capabilities
https://www.genxr.co/#products
MIT License
8.48k stars 2.18k forks source link

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

Open Ebipatei opened 1 year ago

Ebipatei commented 1 year ago

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.

OlafenwaMoses commented 1 year ago

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