From 1075488d893f2167737d89549c3f675b0713aa5a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 4 Dec 2021 16:28:40 +0100 Subject: [PATCH] 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 --- README.md | 2 +- export.py | 23 ++++++++++++----------- 2 files changed, 13 insertions(+), 12 deletions(-) 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) diff --git a/export.py b/export.py index 21c83c697b4d..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 @@ -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__":