Open Elreniel opened 2 years ago
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
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?
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?
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
@thegendolz Hi, what activation function did you use in .pt model? I use silu and it not supported in ONNX.
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?