glenn-jocher
commited on
Commit
β’
1075488
1
Parent(s):
7bf04d9
Single-command multiple-model export (#5882)
Browse files* 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
CHANGED
@@ -148,7 +148,7 @@ $ python train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size
|
|
148 |
* [Roboflow for Datasets, Labeling, and Active Learning](https://github.com/ultralytics/yolov5/issues/4975) π NEW
|
149 |
* [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475)
|
150 |
* [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36) β NEW
|
151 |
-
* [
|
152 |
* [Test-Time Augmentation (TTA)](https://github.com/ultralytics/yolov5/issues/303)
|
153 |
* [Model Ensembling](https://github.com/ultralytics/yolov5/issues/318)
|
154 |
* [Model Pruning/Sparsity](https://github.com/ultralytics/yolov5/issues/304)
|
|
|
148 |
* [Roboflow for Datasets, Labeling, and Active Learning](https://github.com/ultralytics/yolov5/issues/4975) π NEW
|
149 |
* [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475)
|
150 |
* [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36) β NEW
|
151 |
+
* [TFLite, ONNX, CoreML, TensorRT Export](https://github.com/ultralytics/yolov5/issues/251) π
|
152 |
* [Test-Time Augmentation (TTA)](https://github.com/ultralytics/yolov5/issues/303)
|
153 |
* [Model Ensembling](https://github.com/ultralytics/yolov5/issues/318)
|
154 |
* [Model Pruning/Sparsity](https://github.com/ultralytics/yolov5/issues/304)
|
export.py
CHANGED
@@ -2,17 +2,17 @@
|
|
2 |
"""
|
3 |
Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
|
4 |
|
5 |
-
Format | Example |
|
6 |
--- | --- | ---
|
7 |
PyTorch | yolov5s.pt | -
|
8 |
-
TorchScript | yolov5s.torchscript |
|
9 |
-
ONNX | yolov5s.onnx |
|
10 |
-
CoreML | yolov5s.mlmodel |
|
11 |
-
TensorFlow SavedModel | yolov5s_saved_model/ |
|
12 |
-
TensorFlow GraphDef | yolov5s.pb |
|
13 |
-
TensorFlow Lite | yolov5s.tflite |
|
14 |
-
TensorFlow.js | yolov5s_web_model/ |
|
15 |
-
TensorRT | yolov5s.engine |
|
16 |
|
17 |
Usage:
|
18 |
$ 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'
|
|
400 |
def parse_opt():
|
401 |
parser = argparse.ArgumentParser()
|
402 |
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
|
403 |
-
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='
|
404 |
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
|
405 |
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
406 |
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():
|
|
427 |
|
428 |
|
429 |
def main(opt):
|
430 |
-
|
|
|
431 |
|
432 |
|
433 |
if __name__ == "__main__":
|
|
|
2 |
"""
|
3 |
Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
|
4 |
|
5 |
+
Format | Example | `--include ...` argument
|
6 |
--- | --- | ---
|
7 |
PyTorch | yolov5s.pt | -
|
8 |
+
TorchScript | yolov5s.torchscript | `torchscript`
|
9 |
+
ONNX | yolov5s.onnx | `onnx`
|
10 |
+
CoreML | yolov5s.mlmodel | `coreml`
|
11 |
+
TensorFlow SavedModel | yolov5s_saved_model/ | `saved_model`
|
12 |
+
TensorFlow GraphDef | yolov5s.pb | `pb`
|
13 |
+
TensorFlow Lite | yolov5s.tflite | `tflite`
|
14 |
+
TensorFlow.js | yolov5s_web_model/ | `tfjs`
|
15 |
+
TensorRT | yolov5s.engine | `engine`
|
16 |
|
17 |
Usage:
|
18 |
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs
|
|
|
400 |
def parse_opt():
|
401 |
parser = argparse.ArgumentParser()
|
402 |
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
|
403 |
+
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
|
404 |
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
|
405 |
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
406 |
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
|
|
427 |
|
428 |
|
429 |
def main(opt):
|
430 |
+
for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
|
431 |
+
run(**vars(opt))
|
432 |
|
433 |
|
434 |
if __name__ == "__main__":
|