|
|
@@ -30,7 +30,7 @@ if __name__ == '__main__': |
|
|
|
|
|
|
|
# TorchScript export |
|
|
|
try: |
|
|
|
print('\nStarting TorchScript export with torch %s...' % torch.__version) |
|
|
|
print('\nStarting TorchScript export with torch %s...' % torch.__version__) |
|
|
|
f = opt.weights.replace('.pt', '.torchscript') # filename |
|
|
|
ts = torch.jit.trace(model, img) |
|
|
|
ts.save(f) |
|
|
@@ -52,6 +52,9 @@ if __name__ == '__main__': |
|
|
|
onnx_model = onnx.load(f) # load onnx model |
|
|
|
onnx.checker.check_model(onnx_model) # check onnx model |
|
|
|
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' % f) |
|
|
|
except Exception as e: |
|
|
|
print('ONNX export failed: %s' % e) |
|
|
|
|
|
|
|
# Finish |
|
|
|
print('\nExports complete. Visualize with https://github.com/lutzroeder/netron.') |