glenn-jocher commited on
Commit
a9d20eb
1 Parent(s): 569fbc6

export.py update

Browse files
Files changed (1) hide show
  1. 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.__version)
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\nView with https://github.com/lutzroeder/netron' % f)
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.')