TFLite, ONNX, CoreML, TensorRT 내보내기

📚 본 가이드는 훈련된 YOLOv5 🚀 모델을 PyTorch에서 ONNX, TensorRT, CoreML 등을 포함한 다양한 배포 형식으로 내보내는 방법을 설명합니다.

시작하기 전에

Clone repo and install requirements.txt in a Python>=3.8.0 environment, including PyTorch>=1.8. Models and datasets download automatically from the latest YOLOv5 release.

git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install

TensorRT 내보내기 예제(GPU 필요)는 Colab 노트북 부록 섹션을 참조하십시오. Open In Colab

지원되는 내보내기 형식

YOLOv5 추론은 공식적으로 12가지 형식을 지원합니다:

성능 팁
  • ONNX 또는 OpenVINO로 내보내어 최대 3배의 CPU 속도 향상을 경험하십시오. CPU 벤치마크를 참조하십시오.
  • TensorRT로 내보내어 최대 5배의 GPU 속도 향상을 경험하십시오. GPU 벤치마크를 참조하십시오.
형식export.py --include모델
PyTorch-yolov5s.pt
TorchScripttorchscriptyolov5s.torchscript
ONNXonnxyolov5s.onnx
OpenVINOopenvinoyolov5s_openvino_model/
TensorRTengineyolov5s.engine
CoreMLcoremlyolov5s.mlmodel
TensorFlow SavedModelsaved_modelyolov5s_saved_model/
TensorFlow GraphDefpbyolov5s.pb
TensorFlow Litetfliteyolov5s.tflite
TensorFlow Edge TPUedgetpuyolov5s_edgetpu.tflite
TensorFlow.jstfjsyolov5s_web_model/
PaddlePaddlepaddleyolov5s_paddle_model/

벤치마크

아래 벤치마크는 YOLOv5 튜토리얼 노트북을 사용하여 Colab Pro에서 실행되었습니다 Open In Colab. 재현 방법은 다음과 같습니다:

python benchmarks.py --weights yolov5s.pt --imgsz 640 --device 0

Colab Pro V100 GPU

benchmarks: weights=/content/yolov5/yolov5s.pt, imgsz=640, batch_size=1, data=/content/yolov5/data/coco128.yaml, device=0, half=False, test=False
Checking setup...
YOLOv5 🚀 v6.1-135-g7926afc torch 1.10.0+cu111 CUDA:0 (Tesla V100-SXM2-16GB, 16160MiB)
Setup complete ✅ (8 CPUs, 51.0 GB RAM, 46.7/166.8 GB disk)

Benchmarks complete (458.07s)
                   Format  mAP@0.5:0.95  Inference time (ms)
0                 PyTorch        0.4623                10.19
1             TorchScript        0.4623                 6.85
2                    ONNX        0.4623                14.63
3                OpenVINO           NaN                  NaN
4                TensorRT        0.4617                 1.89
5                  CoreML           NaN                  NaN
6   TensorFlow SavedModel        0.4623                21.28
7     TensorFlow GraphDef        0.4623                21.22
8         TensorFlow Lite           NaN                  NaN
9     TensorFlow Edge TPU           NaN                  NaN
10          TensorFlow.js           NaN                  NaN

Colab Pro CPU

benchmarks: weights=/content/yolov5/yolov5s.pt, imgsz=640, batch_size=1, data=/content/yolov5/data/coco128.yaml, device=cpu, half=False, test=False
Checking setup...
YOLOv5 🚀 v6.1-135-g7926afc torch 1.10.0+cu111 CPU
Setup complete ✅ (8 CPUs, 51.0 GB RAM, 41.5/166.8 GB disk)

Benchmarks complete (241.20s)
                   Format  mAP@0.5:0.95  Inference time (ms)
0                 PyTorch        0.4623               127.61
1             TorchScript        0.4623               131.23
2                    ONNX        0.4623                69.34
3                OpenVINO        0.4623                66.52
4                TensorRT           NaN                  NaN
5                  CoreML           NaN                  NaN
6   TensorFlow SavedModel        0.4623               123.79
7     TensorFlow GraphDef        0.4623               121.57
8         TensorFlow Lite        0.4623               316.61
9     TensorFlow Edge TPU           NaN                  NaN
10          TensorFlow.js           NaN                  NaN

훈련된 YOLOv5 모델 내보내기

이 명령은 사전 훈련된 YOLOv5s 모델을 TorchScript 및 ONNX 형식으로 내보냅니다. yolov5s.pt는 사용 가능한 모델 중 두 번째로 작은 'small' 모델입니다. 다른 옵션으로는 yolov5n.pt, yolov5m.pt, yolov5l.pt, yolov5x.pt가 있으며, P6 대응 모델(yolov5s6.pt 등) 또는 사용자 정의 훈련 체크포인트(runs/exp/weights/best.pt)를 사용할 수도 있습니다. 사용 가능한 모든 모델에 대한 자세한 내용은 당사의 README 를 참조하십시오.

python export.py --weights yolov5s.pt --include torchscript onnx

Add --half to export models at FP16 half precision for smaller file sizes

출력:

export: data=data/coco128.yaml, weights=['yolov5s.pt'], imgsz=[640, 640], batch_size=1, device=cpu, half=False, inplace=False, train=False, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, include=['torchscript', 'onnx']
YOLOv5 🚀 v6.2-104-ge3e5122 Python-3.8.0 torch-1.12.1+cu113 CPU

Downloading https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s.pt to yolov5s.pt...
100% 14.1M/14.1M [00:00<00:00, 274MB/s]

Fusing layers...
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients

PyTorch: starting from yolov5s.pt with output shape (1, 25200, 85) (14.1 MB)

TorchScript: starting export with torch 1.12.1+cu113...
TorchScript: export success ✅ 1.7s, saved as yolov5s.torchscript (28.1 MB)

ONNX: starting export with onnx 1.12.0...
ONNX: export success ✅ 2.3s, saved as yolov5s.onnx (28.0 MB)

Export complete (5.5s)
Results saved to /content/yolov5
Detect:          python detect.py --weights yolov5s.onnx
Validate:        python val.py --weights yolov5s.onnx
PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.onnx')
Visualize:       https://netron.app/

3개의 내보낸 모델은 원래의 PyTorch 모델과 함께 저장됩니다:

YOLO export locations

내보낸 모델 시각화를 위해 Netron Viewer를 권장합니다:

YOLO model visualization

내보낸 모델 사용 예제

detect.py는 내보낸 모델에서 추론을 실행합니다:

python detect.py --weights yolov5s.pt             # PyTorch
python detect.py --weights yolov5s.torchscript    # TorchScript
python detect.py --weights yolov5s.onnx           # ONNX Runtime or OpenCV DNN with dnn=True
python detect.py --weights yolov5s_openvino_model # OpenVINO
python detect.py --weights yolov5s.engine         # TensorRT
python detect.py --weights yolov5s.mlmodel        # CoreML (macOS only)
python detect.py --weights yolov5s_saved_model    # TensorFlow SavedModel
python detect.py --weights yolov5s.pb             # TensorFlow GraphDef
python detect.py --weights yolov5s.tflite         # TensorFlow Lite
python detect.py --weights yolov5s_edgetpu.tflite # TensorFlow Edge TPU
python detect.py --weights yolov5s_paddle_model   # PaddlePaddle

val.py는 내보낸 모델에서 검증을 실행합니다:

python val.py --weights yolov5s.pt             # PyTorch
python val.py --weights yolov5s.torchscript    # TorchScript
python val.py --weights yolov5s.onnx           # ONNX Runtime or OpenCV DNN with dnn=True
python val.py --weights yolov5s_openvino_model # OpenVINO
python val.py --weights yolov5s.engine         # TensorRT
python val.py --weights yolov5s.mlmodel        # CoreML (macOS Only)
python val.py --weights yolov5s_saved_model    # TensorFlow SavedModel
python val.py --weights yolov5s.pb             # TensorFlow GraphDef
python val.py --weights yolov5s.tflite         # TensorFlow Lite
python val.py --weights yolov5s_edgetpu.tflite # TensorFlow Edge TPU
python val.py --weights yolov5s_paddle_model   # PaddlePaddle

내보낸 YOLOv5 모델과 함께 PyTorch Hub를 사용하십시오:

import torch

# Model
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.pt")
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.torchscript")  # TorchScript
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.onnx")  # ONNX Runtime
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s_openvino_model")  # OpenVINO
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.engine")  # TensorRT
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.mlmodel")  # CoreML (macOS Only)
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s_saved_model")  # TensorFlow SavedModel
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.pb")  # TensorFlow GraphDef
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.tflite")  # TensorFlow Lite
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s_edgetpu.tflite")  # TensorFlow Edge TPU
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s_paddle_model")  # PaddlePaddle

# Images
img = "https://ultralytics.com/images/zidane.jpg"  # or file, Path, PIL, OpenCV, numpy, list

# Inference
results = model(img)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

OpenCV DNN 추론

ONNX 모델을 사용한 OpenCV 추론:

python export.py --weights yolov5s.pt --include onnx

python detect.py --weights yolov5s.onnx --dnn # detect
python val.py --weights yolov5s.onnx --dnn    # validate

C++ 추론

내보낸 ONNX 모델을 사용한 YOLOv5 OpenCV DNN C++ 추론 예제:

YOLOv5 OpenVINO C++ 추론 예제:

TensorFlow.js 웹 브라우저 추론

지원되는 환경

Ultralytics는 프로젝트를 빠르게 시작할 수 있도록 CUDA, CUDNN, Python, PyTorch와 같은 필수 종속성이 미리 설치된 다양한 준비된 환경을 제공합니다.

프로젝트 상태

YOLOv5 CI

이 배지는 모든 YOLOv5 GitHub Actions CI(지속적 통합) 테스트가 성공적으로 통과되었음을 나타냅니다. 이러한 CI 테스트는 학습, 검증, 추론, 내보내기벤치마크를 포함한 YOLOv5의 다양한 핵심 기능과 성능을 엄격하게 점검합니다. 테스트는 24시간마다 그리고 새로운 커밋이 있을 때마다 수행되며, macOS, Windows 및 Ubuntu에서 일관되고 안정적인 운영을 보장합니다.

댓글