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

Files changed (2) hide show
  1. README.md +1 -1
  2. export.py +12 -11
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
- * [TorchScript, ONNX, CoreML 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)
 
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 | Export `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,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='weights path')
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
- run(**vars(opt))
 
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__":