From 1c4473abadf0bf76d7abebade19c4e84909bbd03 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 4 Dec 2021 15:02:25 +0100 Subject: [PATCH 1/3] 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 ``` --- export.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/export.py b/export.py index 21c83c697b4d..b3709932d235 100644 --- a/export.py +++ b/export.py @@ -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__": From 46a66f0c8a4b30af0e542c895a28b2db8ff22c7e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 4 Dec 2021 15:41:04 +0100 Subject: [PATCH 2/3] Update export.py --- export.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/export.py b/export.py index b3709932d235..88d03a2c9475 100644 --- 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 From 464d3a0783c94aad187da581109a1bf8ad89e383 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 4 Dec 2021 15:43:05 +0100 Subject: [PATCH 3/3] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 6e72d85da7ee..3074330e5505 100644 --- 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)  🌟 NEW * [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475) * [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36)  ⭐ 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)