model.fuse() fix for export.py (#827)

This commit is contained in:
Glenn Jocher committed 2020-08-24 19:27:54 -07:00
1 parent 5f07782f56
commit a8751e50de
2 files changed
+4 -2

No files matched your search

+3 -1
View File
@@ -28,6 +28,9 @@ if __name__ == '__main__':
attempt_download(opt.weights)
model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
model.eval()
model.fuse()
# Update model
model.model[-1].export = True # set Detect() layer export=True
y = model(img) # dry run
@@ -47,7 +50,6 @@ if __name__ == '__main__':
print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
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=['classes', 'boxes'] if y is None else ['output'])