zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

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
This commit is contained in:
Glenn Jocher 2022-02-01 23:52:50 +01:00 committed by GitHub
parent 842d049e1b
commit 4c40933266
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
export.py
View File

@ -16,6 +16,10 @@ TensorFlow Lite | `tflite` | yolov5s.tflite
TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite
TensorFlow.js | `tfjs` | yolov5s_web_model/ 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: Usage:
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ... $ 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 # Exports
f = [''] * 10 # exported filenames f = [''] * 10 # exported filenames
warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning
if 'torchscript' in include: if 'torchscript' in include:
f[0] = export_torchscript(model, im, file, optimize) f[0] = export_torchscript(model, im, file, optimize)
if 'engine' in include: # TensorRT required before ONNX if 'engine' in include: # TensorRT required before ONNX
@ -509,8 +514,6 @@ def parse_opt():
def main(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]): for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
run(**vars(opt)) run(**vars(opt))