Commit
·
7bd7b2c
1
Parent(s):
af41083
update export.py
Browse files- 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 |
-
|
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
|
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=['
|
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
|
58 |
|
59 |
# Finish
|
60 |
-
print('\
|
|
|
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.')
|