From dda0afa22e938b6641665ebf88ccb3b87e4f2461 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 22 Apr 2020 16:00:20 -0700 Subject: [PATCH] onnx export IO layer names update --- detect.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/detect.py b/detect.py index b94903c1..635d7059 100644 --- 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