Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2020-02-25 22:58:26 -08:00
parent 4f3d07f689
commit b12f1a9abe
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
detect.py
View File

@ -45,11 +45,12 @@ def detect(save_img=False):
if ONNX_EXPORT: if ONNX_EXPORT:
model.fuse() model.fuse()
img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192) img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192)
torch.onnx.export(model, img, 'weights/export.onnx', verbose=False, opset_version=11) f = opt.weights.replace(opt.weights.split('.')[-1], 'onnx') # *.onnx filename
torch.onnx.export(model, img, f, verbose=False, opset_version=11)
# Validate exported model # Validate exported model
import onnx import onnx
model = onnx.load('weights/export.onnx') # Load the ONNX model model = onnx.load(f) # Load the ONNX model
onnx.checker.check_model(model) # Check that the IR is well formed onnx.checker.check_model(model) # Check that the IR is well formed
print(onnx.helper.printable_graph(model.graph)) # Print a human readable representation of the graph print(onnx.helper.printable_graph(model.graph)) # Print a human readable representation of the graph
return return