From b12f1a9abe92288c98bfe7da00118ae81a7805e0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Feb 2020 22:58:26 -0800 Subject: [PATCH] updates --- detect.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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