Cách xuất các model PyTorch không phải YOLO bằng Ultralytics#

Ultralytics cung cấp các tiện ích xuất độc lập dưới ultralytics.utils.export để đóng gói nhiều backend đằng sau một giao diện đồng nhất. Bạn có thể xuất bất kỳ torch.nn.Module nào, bao gồm các mô hình hình ảnh timm, các bộ phân loại và phát hiện torchvision, hoặ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 phải tìm hiểu riêng từng backend.

Triển khai các mô hình PyTorch lên môi trường production thường có nghĩa là phải sử dụng nhiều công cụ xuất khác nhau cho từng mục tiêu: torch.onnx.export cho ONNX, coremltools cho các thiết bị Apple, onnx2tf cho TensorFlow, pnnx cho NCNN, vân vân. Mỗi công cụ có API riêng, các điểm đặc thù về phụ thuộc và quy ước đầu ra khác nhau. Các tiện ích này gom tất cả lại thành một mẫu gọi duy nhất.

Tại sao nên sử dụng Ultralytics để xuất các model không phải YOLO?#

• Một API duy nhất cho 10 định dạng: chỉ cần học một quy ước gọi hàm thay vì cả chục loại.
• Bề mặt tiện ích được chia sẻ: Các trình trợ giúp xuất nằm dưới ultralytics.utils.export, do đó một khi các gói backend đã được cài đặt, bạn có thể giữ nguyên cùng một mẫu gọi trên các định dạng.
• Cùng một luồng code như các bản xuất YOLO: các trình trợ giúp tương tự cung cấp sức mạnh cho mọi bản xuất YOLO của Ultralytics.
• Tích hợp sẵn lượng tử hóa FP16 và INT8 cho các định dạng hỗ trợ tính năng này (OpenVINO, CoreML, MNN, NCNN).
• Hoạt động trên CPU: không yêu cầu 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ỳ laptop nào.

Bắt đầu nhanh#

Đường đi nhanh nhất là xuất hai dòng sang ONNX mà không có mã YOLO nào và không cần thiết lập nào 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ợ#

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

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 đã được huấn luyện sẵn từ timm ở chế độ evaluation:

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 kích thước batch động (dynamic batch size), 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 phụ thuộc bổ sung. 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, quantize=16 cho FP16, hoặc quantize=8 cho việc lượng tử hóa INT8. INT8 bổ sung yêu cầu một đối số 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, classifier_names=None, output_file="resnet18.mlpackage")

Đối với các mô hình classification, hãy truyền một danh sách các tên lớp vào classifier_names để thêm một head 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 sẽ đ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
• 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.82
• onnx>=1.12.0,<2.0.0
• protobuf>=5

Xuất sang TensorFlow Frozen Graph#

Tiếp tục từ quá trình xuất SavedModel ở trên, hãy chuyển đổi keras_model được trả về thành một đồ thị .pb dạng đóng băng (frozen):

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 file param và bin có tên cố định cùng với một wrapper Python:

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

torch2ncnn() kiểm tra ncnn và pnnx trong lần sử dụng đầu tiên.

Xuất sang MNN#

Xuất MNN yêu cầu file ONNX làm đầu vào. Hãy xuất sang ONNX trước, sau đó mới 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ợ quantize=16 cho FP16 và quantize=8 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 model PaddlePaddle và các file tham số:

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

Yêu cầu x2paddle và bản phân phối PaddlePaddle tương thích 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 được lưu bên trong thư mục đầu ra:

resnet18_executorch_model/
└── model.pte

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

Xác minh mô hình đã xuất của bạn#

Sau khi xuất, hãy xác minh sự tương đương về mặt số học với mô hình PyTorch gốc trước khi vận chuyển. Một bài kiểm tra nhanh với ONNXBackend từ ultralytics.nn.backends so sánh các đầu ra và gắn cờ lỗi trace hoặc lượng tử hóa 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}")  # typically ~1e-5, well under 1e-4 for FP32

Đố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 mô hình (thường là x cho các mô hình chung), trong khi torch2onnx mặc định là "images".

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

• Hỗ trợ nhiều đầu vào không đồng đều: torch2onnx và torch2openvino chấp nhận một tuple hoặc danh sách các tensor ví dụ cho các mô hình 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. Cài đặt bằng brew install flatbuffers trên macOS hoặc apt install flatbuffers-compiler trên Ubuntu.
• Không suy luận thông qua Ultralytics: Các mô hình phi-YOLO đã xuất không thể được tải lại thông qua YOLO() để suy luận. 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: Các bản xuất Axelera và Sony IMX500 yêu cầu các thuộc tính mô hình cụ thể của YOLO và không có sẵn cho các mô hình chung.
• Các định dạng dành riêng cho nền tảng: TensorRT yêu cầu GPU NVIDIA. RKNN yêu cầu SDK rknn-toolkit2 (chỉ trên Linux). Edge TPU yêu cầu tệp nhị phân edgetpu_compiler (chỉ trên Linux).

Kết luận#

Các tiện ích này đưa bất kỳ mô hình PyTorch nào từ một torch.nn.Module thông thường thành một artifact ONNX, OpenVINO, CoreML, TensorFlow hoặc mobile-runtime sẵn sàng triển khai thông qua một API duy nhất nhất quán. Chọn định dạng phù hợp với phần cứng mục tiêu của bạn, xác minh sự tương đương về số học so với mô hình gốc, sau đó làm theo hướng dẫn tích hợp tương ứng cho các bước triển khai dành riêng cho runtime.

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