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 backendspull/6512/head
parent
842d049e1b
commit
4c40933266
11
export.py
11
export.py
|
@ -16,6 +16,10 @@ TensorFlow Lite | `tflite` | yolov5s.tflite
|
|||
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__":
|
||||
|
|
Loading…
Reference in New Issue