onnx export IO layer names update

2020-04-22 16:00:20 -07:00 · 2020-04-22 16:00:20 -07:00 · dda0afa22e
parent 748f60baae
commit dda0afa22e
1 changed files with 2 additions and 1 deletions
--- a/detect.py
+++ b/detect.py
@ -45,7 +45,8 @@ def detect(save_img=False):
        model.fuse()
        img = torch.zeros((1, 3) + img_size)  # (1, 3, 320, 192)
        f = opt.weights.replace(opt.weights.split('.')[-1], 'onnx')  # *.onnx filename
-        torch.onnx.export(model, img, f, verbose=False, opset_version=11)
+        torch.onnx.export(model, img, f, verbose=False, opset_version=11,
+                          input_names=['images'], output_names=['classes', 'boxes'])

        # Validate exported model
        import onnx