update export.py

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.')
+    print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.')