From 7bd7b2c090cbeb7d5ff4e678d093c258d34368ee Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 3 Jul 2020 20:05:50 -0700 Subject: [PATCH] update export.py --- models/export.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/models/export.py b/models/export.py index 1ff9526..54d7a12 100644 --- a/models/export.py +++ b/models/export.py @@ -26,7 +26,7 @@ if __name__ == '__main__': model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float() model.eval() model.model[-1].export = True # set Detect() layer export=True - _ = model(img) # dry run + y = model(img) # dry run # TorchScript export try: @@ -36,7 +36,7 @@ if __name__ == '__main__': ts.save(f) print('TorchScript export success, saved as %s' % f) except Exception as e: - print('TorchScript export failed: %s' % e) + print('TorchScript export failure: %s' % e) # ONNX export try: @@ -46,7 +46,7 @@ if __name__ == '__main__': 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=['output']) # output_names=['classes', 'boxes'] + output_names=['classes', 'boxes'] if y is None else ['output']) # Checks onnx_model = onnx.load(f) # load onnx model @@ -54,7 +54,7 @@ if __name__ == '__main__': 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 failed: %s' % e) + print('ONNX export failure: %s' % e) # Finish - print('\nExports complete. Visualize with https://github.com/lutzroeder/netron.') \ No newline at end of file + print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.')