Difference between ONNX and pt model

[Elreniel]

I have trained a yolor-p6 model using my own dataset. Also I have converted into ONNX model. When I use .pt model for inference, model could detect the object succesfully. However, the onnx model could not detect the object. Does anyone know what could be the reason for this issue?

[thegendolz]

The problem is that the export.py function is not working properly. I changed the export function such that it works in the same way as detect.py

I have re-written export.py + added all utils python files in the same directory as export.py and changed the module links without utils i.e: from utils.general import -> from general import . Since this caused issues for me. Then I ran the adjusted code and it worked for me!

I also added the cfg file to the export function, since this is required to load the torch file

python3 export.py --weights (Path to your weight file) --img-size 1280 --batch-size 1 --cfg (Path to your cfg file)

import argparse
from torch_utils import select_device

from models import *
from datasets import *
from general import *

if __name__ == '__main__':
    torch.cuda.empty_cache()
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='./yolov4.pt', help='weights path')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--cfg', type=str, default=1, help='batch size')
    opt = parser.parse_args()
    opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
    cfg, imgsz, weights = \
        opt.cfg, opt.img_size, opt.weights

    # Input
    device = select_device('0')
    img = torch.zeros((opt.batch_size, 3, *opt.img_size), device=device)  # image size(1,3,320,192) iDetection
    # Load PyTorch model
    model = Darknet(cfg, imgsz).cuda()
    model.load_state_dict(torch.load(weights, map_location=device)['model'])
    # model = TempModel()
    # model = torch.load_state_dict(torch.load(opt.weights))
    model.eval()
    # model.model[-1].export = True  # set Detect() layer export=True
    y = model(img)  # dry run

    # ONNX export
    try:
        import onnx

        print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
        f = opt.weights.replace('.pt', '.onnx')  # filename
        model.fuse()  # only for ONNX
        torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
                          output_names=['classes', 'boxes'] if y is None else ['output'])

        # Checks
        onnx_model = onnx.load(f)  # load onnx model
        onnx.checker.check_model(onnx_model)  # check onnx model
        print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model
        print('ONNX export success, saved as %s' % f)
    except Exception as e:
        print('ONNX export failure: %s' % e)

Hope this helps you

[Elreniel]

Hey thanks for reply. When I run your code I got the following error;

model = Darknet(cfg, imgsz).cuda()
NameError: name 'Darknet' is not defined

Where did you get the Darknet definition?

[Elreniel]

Hey thanks for reply. When I run your code I got the following error;

model = Darknet(cfg, imgsz).cuda()
NameError: name 'Darknet' is not defined

Where did you get the Darknet definition?

You were talking about main branch. I got the results now. However, the model is working slower than I expected. Can you also share the inference code?

[thegendolz]

What kind of result did you get? How much slower is it? I have tested it and I got almost similar result in comparison to YOLOR detect method

For inference you can view this page: https://docs.microsoft.com/en-us/azure/machine-learning/how-to-inference-onnx-automl-image-models?tabs=object-detect-yolo

[huyhoangle86]

@thegendolz Hi, what activation function did you use in .pt model? I use silu and it not supported in ONNX.