Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

onnx export IO layer names update

This commit is contained in:
Glenn Jocher 2020-04-22 16:00:20 -07:00
parent 748f60baae
commit dda0afa22e
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
detect.py
View File

@ -45,7 +45,8 @@ def detect(save_img=False):
model.fuse() model.fuse()
img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192) img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192)
f = opt.weights.replace(opt.weights.split('.')[-1], 'onnx') # *.onnx filename 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 # Validate exported model
import onnx import onnx