From d45e349a1e65b2d3877e1339497ac0549cdecd4a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 11 Oct 2020 16:23:36 +0200 Subject: [PATCH] Rearrange export input after checks (#1118) img size checks are warnings rather than errors, so current implementation allows improperly formed model inputs. --- models/export.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/models/export.py b/models/export.py index 5eb885d..c5e96f1 100644 --- a/models/export.py +++ b/models/export.py @@ -29,9 +29,6 @@ if __name__ == '__main__': set_logging() t = time.time() - # Input - img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection - # Load PyTorch model model = attempt_load(opt.weights, map_location=torch.device('cpu')) # load FP32 model labels = model.names @@ -40,6 +37,9 @@ if __name__ == '__main__': gs = int(max(model.stride)) # grid size (max stride) 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 + # Update model for k, m in model.named_modules(): m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility