glenn-jocher
commited on
Commit
•
a9d20eb
1
Parent(s):
569fbc6
export.py update
Browse files- models/export.py +5 -2
models/export.py
CHANGED
@@ -30,7 +30,7 @@ if __name__ == '__main__':
|
|
30 |
|
31 |
# TorchScript export
|
32 |
try:
|
33 |
-
print('\nStarting TorchScript export with torch %s...' % torch.
|
34 |
f = opt.weights.replace('.pt', '.torchscript') # filename
|
35 |
ts = torch.jit.trace(model, img)
|
36 |
ts.save(f)
|
@@ -52,6 +52,9 @@ if __name__ == '__main__':
|
|
52 |
onnx_model = onnx.load(f) # load onnx model
|
53 |
onnx.checker.check_model(onnx_model) # check onnx model
|
54 |
print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model
|
55 |
-
print('ONNX export success, saved as %s
|
56 |
except Exception as e:
|
57 |
print('ONNX export failed: %s' % e)
|
|
|
|
|
|
|
|
30 |
|
31 |
# TorchScript export
|
32 |
try:
|
33 |
+
print('\nStarting TorchScript export with torch %s...' % torch.__version__)
|
34 |
f = opt.weights.replace('.pt', '.torchscript') # filename
|
35 |
ts = torch.jit.trace(model, img)
|
36 |
ts.save(f)
|
|
|
52 |
onnx_model = onnx.load(f) # load onnx model
|
53 |
onnx.checker.check_model(onnx_model) # check onnx model
|
54 |
print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model
|
55 |
+
print('ONNX export success, saved as %s' % f)
|
56 |
except Exception as e:
|
57 |
print('ONNX export failed: %s' % e)
|
58 |
+
|
59 |
+
# Finish
|
60 |
+
print('\nExports complete. Visualize with https://github.com/lutzroeder/netron.')
|