TFLite, ONNX, CoreML, TensorRT 내보내기

이 가이드에서는 학습된 YOLOv5 🚀 모델을 PyTorch 에서 ONNX 및 TorchScript 형식으로 내보내는 방법을 설명합니다.

시작하기 전에

리포지토리를 복제하고 요구사항.txt를 설치합니다. Python>=3.8.0 환경을 포함하여 PyTorch>=1.8. 모델 및 데이터 세트는 최신 YOLOv5 릴리스에서 자동으로 다운로드됩니다.

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

For TensorRT 내보내기 예제(GPU 필요)는 Colab 노트북 부록 섹션에 있습니다.

형식

YOLOv5 추론은 공식적으로 11개 형식이 지원됩니다:

💡 프로 팁: ONNX 또는 OpenVINO 으로 내보내면 CPU 속도를 최대 3배까지 높일 수 있습니다. CPU 벤치마크를 참조하세요. 💡 프로팁: TensorRT 으로 내보내면 GPU 속도가 최대 5배 빨라집니다. GPU 벤치마크를 참조하세요.

형식	`export.py --include`	모델
PyTorch	-	`yolov5s.pt`
TorchScript	`torchscript`	`yolov5s.torchscript`
ONNX	`onnx`	`yolov5s.onnx`
OpenVINO	`openvino`	`yolov5s_openvino_model/`
TensorRT	`engine`	`yolov5s.engine`
CoreML	`coreml`	`yolov5s.mlmodel`
TensorFlow SavedModel	`saved_model`	`yolov5s_saved_model/`
TensorFlow GraphDef	`pb`	`yolov5s.pb`
TensorFlow Lite	`tflite`	`yolov5s.tflite`
TensorFlow Edge TPU	`edgetpu`	`yolov5s_edgetpu.tflite`
TensorFlow.js	`tfjs`	`yolov5s_web_model/`
PaddlePaddle	`paddle`	`yolov5s_paddle_model/`

벤치마크

아래 벤치마크는 Colab Pro와 YOLOv5 튜토리얼 노트북에서 실행됩니다. . 재생산하려면:

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 은 두 번째로 작은 모델인 '스몰' 모델입니다. 다른 옵션은 다음과 같습니다. 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

💡 프로 팁: 추가 --half 를 사용하여 더 작은 파일 크기를 위해 FP16 반정밀로 모델을 내보낼 수 있습니다.

출력:

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 내보내기 위치

내보낸 모델을 시각화할 때는 네트론 뷰어를 사용하는 것이 좋습니다:

YOLO 모델 시각화

익스포트된 모델 사용 예시

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

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

val.py 내보낸 모델에 대한 유효성 검사를 실행합니다:

python val.py --weights yolov5s.pt                 # PyTorch
                        yolov5s.torchscript        # TorchScript
                        yolov5s.onnx               # ONNX Runtime or OpenCV DNN with dnn=True
                        yolov5s_openvino_model     # OpenVINO
                        yolov5s.engine             # TensorRT
                        yolov5s.mlmodel            # CoreML (macOS Only)
                        yolov5s_saved_model        # TensorFlow SavedModel
                        yolov5s.pb                 # TensorFlow GraphDef
                        yolov5s.tflite             # TensorFlow Lite
                        yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                        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++ 추론

YOLOv5 내보낸 ONNX 모델 예제에 대한 OpenCV DNN C++ 추론:

YOLOv5 OpenVINO C++ 추론 예제:

TensorFlow.js 웹 브라우저 추론

https://aukerul-shuvo.github.io/YOLOv5_TensorFlow-JS/

지원 환경

Ultralytics 는 바로 사용할 수 있는 다양한 환경을 제공하며, 각 환경에는 CUDA, CUDNN과 같은 필수 종속성이 사전 설치되어 있습니다, Python및 PyTorch와 같은 필수 종속 요소가 사전 설치되어 있어 프로젝트를 시작할 수 있습니다.

무료 GPU 노트북:
Google Cloud: GCP 빠른 시작 가이드
Amazon: AWS 빠른 시작 가이드
Azure: AzureML 빠른 시작 가이드
Docker: Docker 빠른 시작 가이드

프로젝트 상태

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

생성 2023-11-12, 업데이트 2024-01-12
작성자: glenn-jocher (5)