Skip to content

Commit 4c40933

Browse files
authored
Suppress export.run() TracerWarnings (#6499)
Suppresses warnings when calling export.run() directly, not just CLI python export.py. Also adds Requirements examples for CPU and GPU backends
1 parent 842d049 commit 4c40933

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

export.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@
1616
TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite
1717
TensorFlow.js | `tfjs` | yolov5s_web_model/
1818
19+
Requirements:
20+
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU
21+
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU
22+
1923
Usage:
2024
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
2125
@@ -437,6 +441,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
437441

438442
# Exports
439443
f = [''] * 10 # exported filenames
444+
warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning
440445
if 'torchscript' in include:
441446
f[0] = export_torchscript(model, im, file, optimize)
442447
if 'engine' in include: # TensorRT required before ONNX
@@ -509,10 +514,8 @@ def parse_opt():
509514

510515

511516
def main(opt):
512-
with warnings.catch_warnings():
513-
warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning
514-
for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
515-
run(**vars(opt))
517+
for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
518+
run(**vars(opt))
516519

517520

518521
if __name__ == "__main__":

0 commit comments

Comments
 (0)