glenn-jocher commited on
Commit
4c40933
1 Parent(s): 842d049

Suppress export.run() TracerWarnings (#6499)

Browse files

Suppresses warnings when calling export.run() directly, not just CLI python export.py.

Also adds Requirements examples for CPU and GPU backends

Files changed (1) hide show
  1. export.py +7 -4
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
- with warnings.catch_warnings():
513
- warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning
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__":