From 92f5b8bbf0b32f90e698256ed222c56ab4d59739 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 1 Feb 2022 23:52:50 +0100 Subject: [PATCH] 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 --- export.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/export.py b/export.py index 8666f3de63e0..09c50baa415a 100644 --- a/export.py +++ b/export.py @@ -16,6 +16,10 @@ TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite TensorFlow.js | `tfjs` | yolov5s_web_model/ +Requirements: + $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU + $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU + Usage: $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ... @@ -437,6 +441,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' # Exports f = [''] * 10 # exported filenames + warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning if 'torchscript' in include: f[0] = export_torchscript(model, im, file, optimize) if 'engine' in include: # TensorRT required before ONNX @@ -509,10 +514,8 @@ def parse_opt(): def main(opt): - with warnings.catch_warnings(): - warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning - for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]): - run(**vars(opt)) + for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]): + run(**vars(opt)) if __name__ == "__main__":