Export TFLite, ONNX, CoreML, TensorRT#
📚 Hướng dẫn này giải thích cách export model YOLOv5 🚀 đã huấn luyện từ PyTorch sang nhiều định dạng triển khai, bao gồm ONNX, TensorRT, CoreML và các định dạng khác.
Trước khi bắt đầu#
Clone repository và cài đặt requirements.txt trong môi trường Python>=3.8.0, bao gồm PyTorch>=1.8. Model và dataset được tự động tải xuống từ bản phát hành YOLOv5 mới nhất.
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # installXem phần phụ lục ví dụ export TensorRT (yêu cầu GPU) trong notebook Colab của chúng tôi. 
Các định dạng export được hỗ trợ#
Suy luận YOLOv5 được hỗ trợ chính thức ở 12 định dạng:
- Export sang ONNX hoặc OpenVINO để tăng tốc CPU lên đến 3 lần. Xem benchmark CPU.
- Export sang TensorRT để tăng tốc GPU lên đến 5 lần. Xem benchmark GPU.
| Định dạng | export.py --include | Model |
|---|---|---|
| PyTorch | - | yolov5s.pt |
| TorchScript | torchscript | yolov5s.torchscript |
| ONNX | onnx | yolov5s.onnx |
| OpenVINO | openvino | yolov5s_openvino_model/ |
| TensorRT | engine | yolov5s.engine |
| CoreML | coreml | yolov5s.mlpackage |
| 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/ |
Benchmark#
Các benchmark bên dưới chạy trên Colab Pro bằng notebook hướng dẫn YOLOv5 . Để tái hiện:
python benchmarks.py --weights yolov5s.pt --imgsz 640 --device 0GPU V100 trên Colab Pro#
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 NaNCPU trên Colab Pro#
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 NaNExport model YOLOv5 đã huấn luyện#
Lệnh này export model YOLOv5s đã huấn luyện trước sang định dạng TorchScript và ONNX. yolov5s.pt là model "small", nhỏ thứ hai trong số các model hiện có. Các tùy chọn khác gồm yolov5n.pt, yolov5m.pt, yolov5l.pt và yolov5x.pt, cùng các biến thể P6 tương ứng, chẳng hạn yolov5s6.pt, hoặc checkpoint huấn luyện tùy chỉnh của riêng bạn, chẳng hạn runs/train/exp/weights/best.pt. Xem bảng README của chúng tôi để biết chi tiết về tất cả model hiện có.
python export.py --weights yolov5s.pt --include torchscript onnxThêm --half để export model ở độ chính xác bán phần FP16 nhằm giảm kích thước file.
Đầu ra:
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/Các model đã export sẽ được lưu cùng thư mục với model PyTorch gốc:
Khuyến nghị dùng Netron Viewer để trực quan hóa các model đã export:
Ví dụ sử dụng model đã export#
detect.py chạy suy luận trên các model đã export:
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
python detect.py --weights yolov5s_openvino_model # OpenVINO
python detect.py --weights yolov5s.engine # TensorRT
python detect.py --weights yolov5s.mlpackage # 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 # PaddlePaddleval.py chạy đánh giá trên các model đã export:
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
python val.py --weights yolov5s_openvino_model # OpenVINO
python val.py --weights yolov5s.engine # TensorRT
python val.py --weights yolov5s.mlpackage # 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 # PaddlePaddleDùng PyTorch Hub với các model YOLOv5 đã export:
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.mlpackage") # CoreML (chỉ dành cho macOS)
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" # hoặc file, Path, PIL, OpenCV, numpy, list
# Inference
results = model(img)
# Results
results.print() # or .show(), .save(), .crop(), .pandas(), etc.Suy luận OpenCV DNN#
Suy luận OpenCV với model ONNX:
python export.py --weights yolov5s.pt --include onnx
python detect.py --weights yolov5s.onnx --dnn # detect
python val.py --weights yolov5s.onnx --dnn # validateSuy luận C++#
Ví dụ suy luận C++ YOLOv5 OpenCV DNN trên model ONNX đã export:
- https://github.com/Hexmagic/ONNX-yolov5/blob/master/src/test.cpp
- https://github.com/doleron/yolov5-opencv-cpp-python
Ví dụ suy luận C++ YOLOv5 OpenVINO:
- https://github.com/dacquaviva/yolov5-openvino-cpp-python
- https://github.com/UNeedCryDear/yolov5-seg-opencv-onnxruntime-cpp
Suy luận trên trình duyệt web bằng TensorFlow.js#
Môi trường được hỗ trợ#
Ultralytics cung cấp nhiều môi trường sẵn sàng sử dụng, mỗi môi trường được cài đặt sẵn các dependency thiết yếu như CUDA, CUDNN, Python và PyTorch, giúp bạn bắt đầu dự án nhanh chóng.
- Notebook GPU miễn phí:
- Google Cloud: Hướng dẫn bắt đầu nhanh với GCP
- Amazon: Hướng dẫn bắt đầu nhanh với AWS
- Azure: Hướng dẫn bắt đầu nhanh với AzureML
- Docker: Hướng dẫn bắt đầu nhanh với Docker
Trạng thái dự án#
Huy hiệu này cho biết tất cả bài kiểm thử Continuous Integration (CI) của GitHub Actions cho YOLOv5 đều đã chạy thành công. Các bài kiểm thử CI này kiểm tra nghiêm ngặt chức năng và hiệu năng của YOLOv5 trên nhiều khía cạnh chính: huấn luyện, validation, inference, export và benchmark. Các bài kiểm thử đảm bảo hoạt động nhất quán và đáng tin cậy trên macOS, Windows và Ubuntu; chúng được thực hiện mỗi 24 giờ và sau mỗi commit mới.
Câu hỏi thường gặp#
Dùng ONNX hoặc OpenVINO để triển khai trên CPU, nơi chúng chạy nhanh hơn PyTorch đến 3 lần; dùng TensorRT cho GPU NVIDIA, nơi TensorRT chạy nhanh hơn đến 5 lần. CoreML dành cho thiết bị Apple, còn TensorFlow Lite hoặc Edge TPU dành cho Android và phần cứng Coral.
Thêm
--half. Độ chính xác bán phần thường cần GPU đối với hầu hết định dạng, vì vậy hãy truyền thêm--device 0; CoreML là ngoại lệ và có thể export FP16 trên CPU.Có. Thêm
--dynamic. Bản export ONNX sẽ có các trục batch, chiều cao và chiều rộng động, đồng thời được trace trên CPU, vì vậy hãy bỏ qua--device 0. Bản export TensorRT giữ nguyên kích thước ảnh và chỉ thay đổi chiều batch: vẫn cần--device 0, còn--batch-sizeđặt kích thước batch lớn nhất mà engine chấp nhận.Truyền file đã export cho
detect.py --weights,val.py --weightshoặctorch.hub.load("ultralytics/yolov5", "custom", "yolov5s.onnx"). YOLOv5 chọn backend phù hợp dựa trên phần mở rộng của file.