Browse Source

update export.py

5.0
Glenn Jocher 4 years ago
parent
commit
7bd7b2c090
1 changed files with 5 additions and 5 deletions
  1. +5
    -5
      models/export.py

+ 5
- 5
models/export.py View File

model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float() model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
model.eval() model.eval()
model.model[-1].export = True # set Detect() layer export=True model.model[-1].export = True # set Detect() layer export=True
_ = model(img) # dry run
y = model(img) # dry run


# TorchScript export # TorchScript export
try: try:
ts.save(f) ts.save(f)
print('TorchScript export success, saved as %s' % f) print('TorchScript export success, saved as %s' % f)
except Exception as e: except Exception as e:
print('TorchScript export failed: %s' % e)
print('TorchScript export failure: %s' % e)


# ONNX export # ONNX export
try: try:
f = opt.weights.replace('.pt', '.onnx') # filename f = opt.weights.replace('.pt', '.onnx') # filename
model.fuse() # only for ONNX model.fuse() # only for ONNX
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
output_names=['output']) # output_names=['classes', 'boxes']
output_names=['classes', 'boxes'] if y is None else ['output'])


# Checks # Checks
onnx_model = onnx.load(f) # load onnx model onnx_model = onnx.load(f) # load 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' % 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 failure: %s' % e)


# Finish # Finish
print('\nExports complete. Visualize with https://github.com/lutzroeder/netron.')
print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.')

Loading…
Cancel
Save