Suppress export.run() TracerWarnings (#6499)
Browse filesSuppresses warnings when calling export.run() directly, not just CLI python export.py.
Also adds Requirements examples for CPU and GPU backends
export.py
CHANGED
@@ -16,6 +16,10 @@ TensorFlow Lite | `tflite` | yolov5s.tflite
|
|
16 |
TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite
|
17 |
TensorFlow.js | `tfjs` | yolov5s_web_model/
|
18 |
|
|
|
|
|
|
|
|
|
19 |
Usage:
|
20 |
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
|
21 |
|
@@ -437,6 +441,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
|
437 |
|
438 |
# Exports
|
439 |
f = [''] * 10 # exported filenames
|
|
|
440 |
if 'torchscript' in include:
|
441 |
f[0] = export_torchscript(model, im, file, optimize)
|
442 |
if 'engine' in include: # TensorRT required before ONNX
|
@@ -509,10 +514,8 @@ def parse_opt():
|
|
509 |
|
510 |
|
511 |
def main(opt):
|
512 |
-
|
513 |
-
|
514 |
-
for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
|
515 |
-
run(**vars(opt))
|
516 |
|
517 |
|
518 |
if __name__ == "__main__":
|
|
|
16 |
TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite
|
17 |
TensorFlow.js | `tfjs` | yolov5s_web_model/
|
18 |
|
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 |
+
|
23 |
Usage:
|
24 |
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
|
25 |
|
|
|
441 |
|
442 |
# Exports
|
443 |
f = [''] * 10 # exported filenames
|
444 |
+
warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning
|
445 |
if 'torchscript' in include:
|
446 |
f[0] = export_torchscript(model, im, file, optimize)
|
447 |
if 'engine' in include: # TensorRT required before ONNX
|
|
|
514 |
|
515 |
|
516 |
def main(opt):
|
517 |
+
for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
|
518 |
+
run(**vars(opt))
|
|
|
|
|
519 |
|
520 |
|
521 |
if __name__ == "__main__":
|