Browse Source

update export.py

5.0
Glenn Jocher 4 years ago
parent
commit
f1d67f4110
1 changed files with 8 additions and 8 deletions
  1. +8
    -8
      models/export.py

+ 8
- 8
models/export.py View File

"""Exports a YOLOv5 *.pt model to *.onnx and *.torchscript formats
"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats


Usage: Usage:
$ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1 $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
model.model[-1].export = True # set Detect() layer export=True model.model[-1].export = True # set Detect() layer export=True
_ = model(img) # dry run _ = model(img) # dry run


# Export to torchscript
# Export to TorchScript
try: try:
f = opt.weights.replace('.pt', '.torchscript') # filename f = opt.weights.replace('.pt', '.torchscript') # filename
ts = torch.jit.trace(model, img) ts = torch.jit.trace(model, img)
ts.save(f) ts.save(f)
print('Torchscript export success, saved as %s' % f)
except:
print('Torchscript export failed.')
print('TorchScript export success, saved as %s' % f)
except Exception as e:
print('TorchScript export failed: %s' % e)


# Export to ONNX # Export to ONNX
try: try:
f = opt.weights.replace('.pt', '.onnx') # filename f = opt.weights.replace('.pt', '.onnx') # filename
model.fuse() # only for ONNX model.fuse() # only for ONNX
torch.onnx.export(model, img, f, verbose=False, opset_version=11, input_names=['images'],
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
output_names=['output']) # output_names=['classes', 'boxes'] output_names=['output']) # output_names=['classes', 'boxes']


# Checks # Checks
onnx.checker.check_model(onnx_model) # check onnx model onnx.checker.check_model(onnx_model) # check onnx model
print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable representation of the graph print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable representation of the graph
print('ONNX export success, saved as %s\nView with https://github.com/lutzroeder/netron' % f) print('ONNX export success, saved as %s\nView with https://github.com/lutzroeder/netron' % f)
except:
print('ONNX export failed.')
except Exception as e:
print('ONNX export failed: %s' % e)

Loading…
Cancel
Save