Переглянути джерело

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
modifyDataloader
Glenn Jocher GitHub 2 роки тому
джерело
коміт
1075488d89
Не вдалося знайти GPG ключ що відповідає даному підпису Ідентифікатор GPG ключа: 4AEE18F83AFDEB23
2 змінених файлів з 13 додано та 12 видалено
  1. +1
    -1
      README.md
  2. +12
    -11
      export.py

+ 1
- 1
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)

+ 12
- 11
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__":

Завантаження…
Відмінити
Зберегти