glenn-jocher commited on
Commit
7bd7b2c
·
1 Parent(s): af41083

update export.py

Browse files
Files changed (1) hide show
  1. models/export.py +5 -5
models/export.py CHANGED
@@ -26,7 +26,7 @@ if __name__ == '__main__':
26
  model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
27
  model.eval()
28
  model.model[-1].export = True # set Detect() layer export=True
29
- _ = model(img) # dry run
30
 
31
  # TorchScript export
32
  try:
@@ -36,7 +36,7 @@ if __name__ == '__main__':
36
  ts.save(f)
37
  print('TorchScript export success, saved as %s' % f)
38
  except Exception as e:
39
- print('TorchScript export failed: %s' % e)
40
 
41
  # ONNX export
42
  try:
@@ -46,7 +46,7 @@ if __name__ == '__main__':
46
  f = opt.weights.replace('.pt', '.onnx') # filename
47
  model.fuse() # only for ONNX
48
  torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
49
- output_names=['output']) # output_names=['classes', 'boxes']
50
 
51
  # Checks
52
  onnx_model = onnx.load(f) # load onnx model
@@ -54,7 +54,7 @@ if __name__ == '__main__':
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.')
 
26
  model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
27
  model.eval()
28
  model.model[-1].export = True # set Detect() layer export=True
29
+ y = model(img) # dry run
30
 
31
  # TorchScript export
32
  try:
 
36
  ts.save(f)
37
  print('TorchScript export success, saved as %s' % f)
38
  except Exception as e:
39
+ print('TorchScript export failure: %s' % e)
40
 
41
  # ONNX export
42
  try:
 
46
  f = opt.weights.replace('.pt', '.onnx') # filename
47
  model.fuse() # only for ONNX
48
  torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
49
+ output_names=['classes', 'boxes'] if y is None else ['output'])
50
 
51
  # Checks
52
  onnx_model = onnx.load(f) # load 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 failure: %s' % e)
58
 
59
  # Finish
60
+ print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.')