Cách xuất mô hình PyTorch không phải YOLO với Ultralytics

Triển khai mô hình PyTorch lên production thường đòi hỏi phải sử dụng các bộ xuất (exporter) khác nhau cho mỗi mục tiêu: torch.onnx.export cho ONNX, coremltools cho các thiết bị Apple, onnx2tf cho TensorFlow, pnnx cho NCNN, v.v. Mỗi công cụ đều có API riêng, các phụ thuộc phức tạp và quy ước đầu ra khác nhau.

Ultralytics cung cấp các tiện ích xuất độc lập bao bọc nhiều backend dưới một giao diện nhất quán. Bạn có thể xuất bất kỳ mô hình hình ảnh torch.nn.Module, bao gồm timm nào, các bộ phân loại và phát hiện torchvision hoặc các kiến trúc tùy chỉnh của riêng bạn sang ONNX, TorchScript, OpenVINO, CoreML, NCNN, PaddlePaddle, MNN, ExecuTorch, và TensorFlow SavedModel mà không cần học riêng từng backend.

Tại sao nên sử dụng Ultralytics để xuất mô hình không phải YOLO?

• Một API cho 10 định dạng: chỉ cần học một quy ước gọi hàm thay vì cả tá cái khác.
• Bề mặt tiện ích chung: các trình hỗ trợ xuất nằm dưới ultralytics.utils.export, vì vậy khi đã cài đặt các gói backend, bạn có thể giữ nguyên mẫu gọi hàm trên khắp các định dạng.
• Cùng đường dẫn mã nguồn như khi xuất YOLO: cùng các trình hỗ trợ đó cung cấp sức mạnh cho mọi tác vụ xuất YOLO của Ultralytics.
• Lượng tử hóa FP16 và INT8 được tích hợp sẵn cho các định dạng hỗ trợ (OpenVINO, CoreML, MNN, NCNN).
• Hoạt động trên CPU: không cần GPU cho chính bước xuất, vì vậy bạn có thể chạy nó cục bộ trên bất kỳ máy tính xách tay nào.

Bắt đầu nhanh

Con đường nhanh nhất là xuất hai dòng sang ONNX mà không cần mã YOLO và không cần thiết lập gì ngoài pip install ultralytics onnx timm:

import timm
import torch

from ultralytics.utils.export import torch2onnx

model = timm.create_model("resnet18", pretrained=True).eval()
torch2onnx(model, torch.randn(1, 3, 224, 224), output_file="resnet18.onnx")

Các định dạng xuất được hỗ trợ

Lớp model torch2* các hàm nhận một torch.nn.Module tiêu chuẩn và một input tensor mẫu. MNN, TF SavedModel và TF Frozen Graph đi qua một artifact ONNX hoặc Keras trung gian. Không yêu cầu bất kỳ thuộc tính cụ thể nào của YOLO trong cả hai trường hợp.

Các ví dụ từng bước

Mọi ví dụ dưới đây đều sử dụng cùng một thiết lập, một ResNet-18 đã huấn luyện trước từ timm ở chế độ đánh giá:

import timm
import torch

model = timm.create_model("resnet18", pretrained=True).eval()
im = torch.randn(1, 3, 224, 224)

Xuất sang ONNX

from ultralytics.utils.export import torch2onnx

torch2onnx(model, im, output_file="resnet18.onnx")

Đối với batch size động, hãy truyền một từ điển dynamic:

torch2onnx(model, im, output_file="resnet18_dyn.onnx", dynamic={"images": {0: "batch_size"}})

Opset mặc định là 14 và tên đầu vào mặc định là "images". Ghi đè bằng các đối số opset, input_names, hoặc output_names.

Xuất sang TorchScript

Không cần thêm phụ thuộc nào. Sử dụng torch.jit.trace bên dưới.

from ultralytics.utils.export import torch2torchscript

torch2torchscript(model, im, output_file="resnet18.torchscript")

Xuất sang OpenVINO

from ultralytics.utils.export import torch2openvino

ov_model = torch2openvino(model, im, output_dir="resnet18_openvino_model")

Thư mục chứa một cặp model.xml và model.bin có tên cố định:

resnet18_openvino_model/
├── model.xml
└── model.bin

Truyền dynamic=True cho các hình dạng đầu vào động, half=True cho FP16, hoặc int8=True cho lượng tử hóa INT8. INT8 yêu cầu thêm một calibration_dataset.

Yêu cầu openvino>=2024.0.0 (hoặc >=2025.2.0 trên macOS 15.4+) và torch>=2.1.

Xuất sang CoreML

import coremltools as ct

from ultralytics.utils.export import torch2coreml

inputs = [ct.TensorType("input", shape=(1, 3, 224, 224))]
ct_model = torch2coreml(model, inputs, im, output_file="resnet18.mlpackage")

Đối với classification các mô hình, hãy truyền danh sách tên lớp vào classifier_names để thêm một phần phân loại vào mô hình CoreML.

Yêu cầu coremltools>=9.0, torch>=1.11, và numpy<=2.3.5. Không được hỗ trợ trên Windows.

Xuất sang TensorFlow SavedModel

Việc xuất TF SavedModel đi qua ONNX như một bước trung gian:

from ultralytics.utils.export import onnx2saved_model, torch2onnx

torch2onnx(model, im, output_file="resnet18.onnx")
keras_model = onnx2saved_model("resnet18.onnx", output_dir="resnet18_saved_model")

Hàm trả về một mô hình Keras và cũng tạo ra các tệp TFLite (.tflite) bên trong thư mục đầu ra:

resnet18_saved_model/
├── saved_model.pb
├── variables/
├── resnet18_float32.tflite
├── resnet18_float16.tflite
└── resnet18_int8.tflite

Yêu cầu:

• tensorflow>=2.0.0,<=2.19.0
• onnx2tf>=1.26.3,<1.29.0
• tf_keras<=2.19.0
• sng4onnx>=1.0.1
• onnx_graphsurgeon>=0.3.26 (cài đặt với --extra-index-url https://pypi.ngc.nvidia.com)
• ai-edge-litert>=1.2.0,<1.4.0 trên macOS (ai-edge-litert>=1.2.0 trên các nền tảng khác)
• onnxslim>=0.1.71
• onnx>=1.12.0,<2.0.0
• protobuf>=5

Xuất sang TensorFlow Frozen Graph

Tiếp tục từ phần xuất SavedModel ở trên, chuyển đổi mô hình Keras trả về thành một đồ thị .pb đóng băng:

from pathlib import Path

from ultralytics.utils.export import keras2pb

keras2pb(keras_model, output_file=Path("resnet18_saved_model/resnet18.pb"))

Xuất sang NCNN

from ultralytics.utils.export import torch2ncnn

torch2ncnn(model, im, output_dir="resnet18_ncnn_model")

Thư mục chứa các tệp param và bin có tên cố định cùng với một trình bao bọc Python:

resnet18_ncnn_model/
├── model.ncnn.param
├── model.ncnn.bin
└── model_ncnn.py

torch2ncnn() kiểm tra ncnn và pnnx mới nhất của Ultralytics trong lần sử dụng đầu tiên.

Xuất sang MNN

Xuất MNN yêu cầu tệp ONNX làm đầu vào. Hãy xuất sang ONNX trước, sau đó chuyển đổi:

from ultralytics.utils.export import onnx2mnn, torch2onnx

torch2onnx(model, im, output_file="resnet18.onnx")
onnx2mnn("resnet18.onnx", output_file="resnet18.mnn")

Hỗ trợ half=True cho FP16 và int8=True cho lượng tử hóa INT8. Yêu cầu MNN>=2.9.6 và torch>=1.10.

Xuất sang PaddlePaddle

from ultralytics.utils.export import torch2paddle

torch2paddle(model, im, output_dir="resnet18_paddle_model")

Thư mục chứa mô hình PaddlePaddle và các tệp tham số:

resnet18_paddle_model/
├── model.pdmodel
└── model.pdiparams

Yêu cầu x2paddle và bản phân phối PaddlePaddle chính xác cho nền tảng của bạn:

• paddlepaddle-gpu>=3.0.0,<3.3.0 trên CUDA
• paddlepaddle==3.0.0 trên CPU ARM64
• paddlepaddle>=3.0.0,<3.3.0 trên các CPU khác

Không được hỗ trợ trên NVIDIA Jetson.

Xuất sang ExecuTorch

from ultralytics.utils.export import torch2executorch

torch2executorch(model, im, output_dir="resnet18_executorch_model")

Tệp .pte được xuất sẽ được lưu trong thư mục đầu ra:

resnet18_executorch_model/
└── model.pte

Yêu cầu torch>=2.9.0 và runtime ExecuTorch tương ứng (pip install executorch). Để biết cách sử dụng runtime, hãy xem tích hợp ExecuTorch.

Xác minh model đã xuất của bạn

Sau khi xuất, hãy xác minh sự tương đương về mặt số liệu với model PyTorch gốc trước khi triển khai. Một bài kiểm tra nhanh với ONNXBackend từ ultralytics.nn.backends sẽ so sánh các kết quả đầu ra và gắn cờ các lỗi tracing hoặc quantization ngay từ sớm:

import numpy as np
import timm
import torch

from ultralytics.nn.backends import ONNXBackend

model = timm.create_model("resnet18", pretrained=True).eval()
im = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    pytorch_output = model(im).numpy()

onnx_model = ONNXBackend("resnet18.onnx", device=torch.device("cpu"))
onnx_output = onnx_model(im)[0]

diff = np.abs(pytorch_output - onnx_output).max()
print(f"Max difference: {diff:.6f}")  # should be < 1e-5

Đối với các runtime khác, tên tensor đầu vào có thể khác. Ví dụ: OpenVINO sử dụng tên đối số forward của model (thường là x cho các model chung), trong khi torch2onnx mặc định là "images".

Các hạn chế đã biết

• Hỗ trợ nhiều đầu vào (multi-input) chưa đồng nhất: torch2onnx và torch2openvino chấp nhận một tuple hoặc danh sách các tensor ví dụ cho các model có nhiều đầu vào. torch2torchscript, torch2coreml, torch2ncnn, torch2paddle, và torch2executorch giả định một tensor đầu vào duy nhất.
• ExecuTorch cần flatc: Runtime ExecuTorch yêu cầu trình biên dịch FlatBuffers. Hãy cài đặt bằng brew install flatbuffers trên macOS hoặc apt install flatbuffers-compiler trên Ubuntu.
• Không hỗ trợ inference thông qua Ultralytics: Các model phi-YOLO đã xuất không thể được tải lại thông qua YOLO() để thực hiện inference. Hãy sử dụng runtime gốc cho từng định dạng (ONNX Runtime, OpenVINO Runtime, v.v.).
• Các định dạng chỉ dành cho YOLO: Axelera và Sony IMX500 các bản xuất yêu cầu các thuộc tính model cụ thể của YOLO và không khả dụng cho các model chung.
• Các định dạng cụ thể theo nền tảng: TensorRT yêu cầu GPU NVIDIA. RKNN yêu cầu rknn-toolkit2 SDK (chỉ dành cho Linux). Edge TPU yêu cầu edgetpu_compiler binary (chỉ dành cho Linux).

Câu hỏi thường gặp

Tôi có thể xuất những model nào với Ultralytics?

Bất kỳ torch.nn.Module. Điều này bao gồm các model từ timm, torchvision, hoặc bất kỳ model PyTorch tùy chỉnh nào. Model phải ở chế độ đánh giá (model.eval()) trước khi xuất. ONNX và OpenVINO cũng chấp nhận một tuple các tensor ví dụ cho các model có nhiều đầu vào.

Định dạng xuất nào hoạt động mà không cần GPU?

Tất cả các định dạng được hỗ trợ (TorchScript, ONNX, OpenVINO, CoreML, TF SavedModel, TF Frozen Graph, NCNN, PaddlePaddle, MNN, ExecuTorch) đều có thể xuất trên CPU. Không cần GPU cho chính quá trình xuất. TensorRT là định dạng duy nhất yêu cầu GPU NVIDIA.

Tôi cần phiên bản Ultralytics nào?

Sử dụng Ultralytics >=8.4.38, bao gồm mô-đun ultralytics.utils.export và output_file/output_dir.

Tôi có thể xuất model torchvision sang CoreML để triển khai trên iOS không?

Có. Các model phân loại, phát hiện và phân đoạn (segmentation) của torchvision có thể xuất sang .mlpackage thông qua torch2coreml. Đối với các model phân loại hình ảnh, hãy truyền một danh sách tên lớp vào classifier_names để tích hợp classification head. Hãy thực hiện xuất trên macOS hoặc Linux. CoreML không được hỗ trợ trên Windows. Xem tích hợp CoreML để biết chi tiết triển khai trên iOS.

Tôi có thể quantize model đã xuất của mình sang INT8 hoặc FP16 không?

Có, đối với một số định dạng. Hãy truyền half=True cho FP16 hoặc int8=True cho INT8 khi xuất sang OpenVINO, CoreML, MNN hoặc NCNN. INT8 trong OpenVINO cũng yêu cầu một đối số calibration_dataset cho post-training quantization. Xem trang tích hợp của từng định dạng để biết các đánh đổi về quantization.