From a62a1c2c679cc6de730debf1529f073d10180452 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 1 Jul 2020 16:14:49 -0700 Subject: [PATCH] export.py update --- detect.py | 8 ++++---- models/export.py | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/detect.py b/detect.py index bb84a0d..93faf6d 100644 --- a/detect.py +++ b/detect.py @@ -156,9 +156,9 @@ if __name__ == '__main__': print(opt) with torch.no_grad(): - detect() + # detect() # 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) + for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov3-spp.pt']: + detect() + create_pretrained(opt.weights, opt.weights) diff --git a/models/export.py b/models/export.py index 1c78d3a..bb310f3 100644 --- a/models/export.py +++ b/models/export.py @@ -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,7 +28,7 @@ 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) @@ -39,8 +37,10 @@ if __name__ == '__main__': 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=12, input_names=['images'],