Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

export.py update

Browse Source
This commit is contained in:
Glenn Jocher committed 2020-07-03 11:29:53 -07:00
1 parent 3bdea3f697
commit 569fbc63c8
1 file changed
+3 -1
Download Patch File Download Diff File Expand all files Collapse all files
models/export.py
+3 -1
View File
@@ -30,6 +30,7 @@ if __name__ == '__main__':
# TorchScript export # TorchScript export
try: try:
print('\nStarting TorchScript export with torch %s...' % torch.__version)
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)
@@ -41,6 +42,7 @@ if __name__ == '__main__':
try: try:
import onnx import onnx
print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
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=12, input_names=['images'], torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
@@ -49,7 +51,7 @@ if __name__ == '__main__':
# Checks # Checks
onnx_model = onnx.load(f) # load onnx model onnx_model = onnx.load(f) # load onnx model
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 model
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 Exception as e: except Exception as e:
print('ONNX export failed: %s' % e) print('ONNX export failed: %s' % e)