diff --git a/detect.py b/detect.py index dbf32028..64a7867d 100644 --- a/detect.py +++ b/detect.py @@ -45,11 +45,12 @@ def detect(save_img=False): if ONNX_EXPORT: model.fuse() 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 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 print(onnx.helper.printable_graph(model.graph)) # Print a human readable representation of the graph return