diff --git a/detect.py b/detect.py index bb84a0d..2650c20 100644 --- a/detect.py +++ b/detect.py @@ -158,7 +158,7 @@ if __name__ == '__main__': with torch.no_grad(): detect() - # Update all models + # # Update all models # for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov3-spp.pt']: - # detect() - # create_pretrained(opt.weights, opt.weights) + # detect() + # create_pretrained(opt.weights, opt.weights) diff --git a/models/export.py b/models/export.py index 2aa6ce4..bb310f3 100644 --- a/models/export.py +++ b/models/export.py @@ -1,4 +1,4 @@ -"""Exports a YOLOv5 *.pt model to *.onnx and *.torchscript formats +"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats Usage: $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1 @@ -6,8 +6,6 @@ Usage: import argparse -import onnx - from models.common import * from utils import google_utils @@ -21,7 +19,7 @@ if __name__ == '__main__': print(opt) # Input - img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size, (1, 3, 320, 192) iDetection + img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection # Load PyTorch model google_utils.attempt_download(opt.weights) @@ -30,20 +28,22 @@ if __name__ == '__main__': model.model[-1].export = True # set Detect() layer export=True _ = model(img) # dry run - # Export to torchscript + # TorchScript export try: f = opt.weights.replace('.pt', '.torchscript') # filename ts = torch.jit.trace(model, img) ts.save(f) - print('Torchscript export success, saved as %s' % f) - except: - print('Torchscript export failed.') + print('TorchScript export success, saved as %s' % f) + except Exception as e: + print('TorchScript export failed: %s' % e) - # Export to ONNX + # ONNX export try: + import onnx + f = opt.weights.replace('.pt', '.onnx') # filename model.fuse() # only for ONNX - torch.onnx.export(model, img, f, verbose=False, opset_version=11, input_names=['images'], + torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], output_names=['output']) # output_names=['classes', 'boxes'] # Checks @@ -51,5 +51,5 @@ if __name__ == '__main__': onnx.checker.check_model(onnx_model) # check onnx model print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable representation of the graph print('ONNX export success, saved as %s\nView with https://github.com/lutzroeder/netron' % f) - except: - print('ONNX export failed.') + except Exception as e: + print('ONNX export failed: %s' % e)