Skip to content

Commit

Permalink
Suppress export.run() TracerWarnings (ultralytics#6499)
Browse files Browse the repository at this point in the history
Suppresses warnings when calling export.run() directly, not just CLI python export.py.

Also adds Requirements examples for CPU and GPU backends
  • Loading branch information
glenn-jocher committed Feb 1, 2022
1 parent 77a3ecf commit 92f5b8b
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions export.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 ...
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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__":
Expand Down

0 comments on commit 92f5b8b

Please sign in to comment.