Glenn Jocher
4 years ago
1 changed files with 5 additions and 2 deletions
+ 5
- 2
models/export.py
View File
| # TorchScript export | # TorchScript export | ||||
| try: | 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 | 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) | ||||
| 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 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: | except Exception as e: | ||||
| print('ONNX export failed: %s' % e) | print('ONNX export failed: %s' % e) | ||||
| # Finish | |||||
| print('\nExports complete. Visualize with https://github.com/lutzroeder/netron.') |
Loading…