|
|
|
|
|
|
|
|
""" |
|
|
""" |
|
|
Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit |
|
|
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 | - |
|
|
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: |
|
|
Usage: |
|
|
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs |
|
|
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs |
|
|
|
|
|
|
|
|
def parse_opt(): |
|
|
def parse_opt(): |
|
|
parser = argparse.ArgumentParser() |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') |
|
|
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('--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('--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') |
|
|
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def main(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__": |
|
|
if __name__ == "__main__": |