From 49bd2fe4a1aabb01333a29028b35ce7dff6b5fda Mon Sep 17 00:00:00 2001 From: Jan Hajek Date: Thu, 25 Feb 2021 11:55:57 +0100 Subject: [PATCH 1/7] option for skip last layer and cuda export support --- models/export.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/models/export.py b/models/export.py index cc817871f218..417ed0721960 100644 --- a/models/export.py +++ b/models/export.py @@ -24,6 +24,7 @@ parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') parser.add_argument('--batch-size', type=int, default=1, help='batch size') + parser.add_argument('--skip-last-layer', action='store_true', help='skip export of last (detect) layer') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand print(opt) @@ -31,7 +32,8 @@ t = time.time() # Load PyTorch model - model = attempt_load(opt.weights, map_location=torch.device('cpu')) # load FP32 model + device = torch.device("cuda:0") if torch.cuda.is_available() else "cpu" + model = attempt_load(opt.weights, map_location=device) # load FP32 model labels = model.names # Checks @@ -39,7 +41,7 @@ opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples # 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).to(device) # image size(1,3,320,192) iDetection # Update model for k, m in model.named_modules(): @@ -51,7 +53,10 @@ m.act = SiLU() # elif isinstance(m, models.yolo.Detect): # m.forward = m.forward_export # assign forward (optional) - model.model[-1].export = True # set Detect() layer export=True + if opt.skip_last_layer: + model.model[-1].export = False # set Detect() layer export=False + else: + model.model[-1].export = True # set Detect() layer export=True y = model(img) # dry run # TorchScript export From 93b41d7d8c1cdf8f29ef74619a4b3da0ba0a7308 Mon Sep 17 00:00:00 2001 From: Jan Hajek Date: Tue, 2 Mar 2021 17:54:05 +0100 Subject: [PATCH 2/7] added parameter device --- models/export.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/models/export.py b/models/export.py index 417ed0721960..90ccc5af034e 100644 --- a/models/export.py +++ b/models/export.py @@ -8,6 +8,8 @@ import sys import time +from utils.torch_utils import select_device + sys.path.append('./') # to run '$ python *.py' files in subdirectories import torch @@ -25,6 +27,7 @@ parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--skip-last-layer', action='store_true', help='skip export of last (detect) layer') + parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand print(opt) @@ -32,7 +35,7 @@ t = time.time() # Load PyTorch model - device = torch.device("cuda:0") if torch.cuda.is_available() else "cpu" + device = select_device(opt.device) model = attempt_load(opt.weights, map_location=device) # load FP32 model labels = model.names From e3f5bdb79ae33f055c09be0041b4495f428c9bcd Mon Sep 17 00:00:00 2001 From: Jan Hajek Date: Tue, 2 Mar 2021 22:00:46 +0100 Subject: [PATCH 3/7] fix import --- models/export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/export.py b/models/export.py index 90ccc5af034e..5d5ddc2e6c3f 100644 --- a/models/export.py +++ b/models/export.py @@ -8,7 +8,6 @@ import sys import time -from utils.torch_utils import select_device sys.path.append('./') # to run '$ python *.py' files in subdirectories @@ -19,6 +18,7 @@ from models.experimental import attempt_load from utils.activations import Hardswish, SiLU from utils.general import set_logging, check_img_size +from utils.torch_utils import select_device if __name__ == '__main__': parser = argparse.ArgumentParser() From 6373c3fd4902af8a729516d837835137cf18bee5 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 6 Mar 2021 10:37:14 -0800 Subject: [PATCH 4/7] cleanup 1 --- models/export.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/models/export.py b/models/export.py index 5d5ddc2e6c3f..8b89ba1c07e7 100644 --- a/models/export.py +++ b/models/export.py @@ -8,7 +8,6 @@ import sys import time - sys.path.append('./') # to run '$ python *.py' files in subdirectories import torch @@ -24,9 +23,9 @@ parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/ parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width - parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') parser.add_argument('--batch-size', type=int, default=1, help='batch size') - parser.add_argument('--skip-last-layer', action='store_true', help='skip export of last (detect) layer') + parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') + parser.add_argument('--skip-grid', action='store_true', help='skip Detect() layer grid computation') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand @@ -56,10 +55,7 @@ m.act = SiLU() # elif isinstance(m, models.yolo.Detect): # m.forward = m.forward_export # assign forward (optional) - if opt.skip_last_layer: - model.model[-1].export = False # set Detect() layer export=False - else: - model.model[-1].export = True # set Detect() layer export=True + model.model[-1].export = not opt.skip_grid # set Detect() layer export Boolean y = model(img) # dry run # TorchScript export From b482563553f4587baa5f86a64be707ab8c9083c3 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 6 Mar 2021 10:56:11 -0800 Subject: [PATCH 5/7] cleanup 2 --- models/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/export.py b/models/export.py index 8b89ba1c07e7..3208f60a5f10 100644 --- a/models/export.py +++ b/models/export.py @@ -25,7 +25,7 @@ parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') - parser.add_argument('--skip-grid', action='store_true', help='skip Detect() layer grid computation') + parser.add_argument('--no-grid', action='store_true', help='no Detect() layer grid computation') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand @@ -55,7 +55,7 @@ m.act = SiLU() # elif isinstance(m, models.yolo.Detect): # m.forward = m.forward_export # assign forward (optional) - model.model[-1].export = not opt.skip_grid # set Detect() layer export Boolean + model.model[-1].export = not opt.no_grid # set Detect() layer grid export y = model(img) # dry run # TorchScript export From 435cb8339dceab316630fed8ff418f6d9f500ce8 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 6 Mar 2021 11:12:47 -0800 Subject: [PATCH 6/7] opt-in grid --grid will export with grid computation, default export will skip grid (same as current) --- models/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/export.py b/models/export.py index 3208f60a5f10..7a9a3c249333 100644 --- a/models/export.py +++ b/models/export.py @@ -25,7 +25,7 @@ parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') - parser.add_argument('--no-grid', action='store_true', help='no Detect() layer grid computation') + parser.add_argument('--grid', action='store_true', help='export Detect() layer grid') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand @@ -55,7 +55,7 @@ m.act = SiLU() # elif isinstance(m, models.yolo.Detect): # m.forward = m.forward_export # assign forward (optional) - model.model[-1].export = not opt.no_grid # set Detect() layer grid export + model.model[-1].export = not opt.grid # set Detect() layer grid export y = model(img) # dry run # TorchScript export From 2b2c800160e59f5109efdf8e4f399a6800a1da7c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 6 Mar 2021 11:16:10 -0800 Subject: [PATCH 7/7] default --device cpu GPU export causes ONNX and CoreML errors. --- models/export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/export.py b/models/export.py index 7a9a3c249333..11e60c7a583d 100644 --- a/models/export.py +++ b/models/export.py @@ -26,7 +26,7 @@ parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') parser.add_argument('--grid', action='store_true', help='export Detect() layer grid') - parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') + parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand print(opt)