|
|
|
|
|
|
|
|
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.') |