Single-command multiple-model export (#5882)

* Export multiple models in series Export multiple models in series by adding additional `*.pt` files to the `--weights` argument, i.e.: ```bash python export.py --include tflite --weights yolov5n.pt # export 1 model python export.py --include tflite --weights yolov5n.pt yolov5s.pt yolov5m.pt yolov5l.pt yolov5x.pt # export 5 models ``` * Update export.py * Update README.md
ultralytics · Dec 4, 2021 · 1075488 · 1075488
1 parent 7bf04d9
commit 1075488
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -148,7 +148,7 @@ $ python train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size
 * [Roboflow for Datasets, Labeling, and Active Learning](https://github.com/ultralytics/yolov5/issues/4975)&nbsp; 🌟 NEW
 * [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475)
 * [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36)&nbsp; ⭐ NEW
-* [TorchScript, ONNX, CoreML Export](https://github.com/ultralytics/yolov5/issues/251) 🚀
+* [TFLite, ONNX, CoreML, TensorRT Export](https://github.com/ultralytics/yolov5/issues/251) 🚀
 * [Test-Time Augmentation (TTA)](https://github.com/ultralytics/yolov5/issues/303)
 * [Model Ensembling](https://github.com/ultralytics/yolov5/issues/318)
 * [Model Pruning/Sparsity](https://github.com/ultralytics/yolov5/issues/304)

diff --git a/export.py b/export.py
@@ -2,17 +2,17 @@
 """
 Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
 
-Format                  | Example                   | Export `include=(...)` argument
+Format                  | Example                   | `--include ...` argument
 ---                     | ---                       | ---
 PyTorch                 | yolov5s.pt                | -
-TorchScript             | yolov5s.torchscript       | 'torchscript'
-ONNX                    | yolov5s.onnx              | 'onnx'
-CoreML                  | yolov5s.mlmodel           | 'coreml'
-TensorFlow SavedModel   | yolov5s_saved_model/      | 'saved_model'
-TensorFlow GraphDef     | yolov5s.pb                | 'pb'
-TensorFlow Lite         | yolov5s.tflite            | 'tflite'
-TensorFlow.js           | yolov5s_web_model/        | 'tfjs'
-TensorRT                | yolov5s.engine            | 'engine'
+TorchScript             | yolov5s.torchscript       | `torchscript`
+ONNX                    | yolov5s.onnx              | `onnx`
+CoreML                  | yolov5s.mlmodel           | `coreml`
+TensorFlow SavedModel   | yolov5s_saved_model/      | `saved_model`
+TensorFlow GraphDef     | yolov5s.pb                | `pb`
+TensorFlow Lite         | yolov5s.tflite            | `tflite`
+TensorFlow.js           | yolov5s_web_model/        | `tfjs`
+TensorRT                | yolov5s.engine            | `engine`
 
 Usage:
     $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs
@@ -400,7 +400,7 @@ def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
 def parse_opt():
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
-    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path')
+    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
     parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
     parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
@@ -427,7 +427,8 @@ def parse_opt():
 
 
 def main(opt):
-    run(**vars(opt))
+    for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
+        run(**vars(opt))
 
 
 if __name__ == "__main__":