Как экспортировать PyTorch-модели (не YOLO) с помощью Ultralytics

Развертывание моделей PyTorch в продакшн обычно означает использование отдельного экспортёра для каждой цели: torch.onnx.export для ONNX, coremltools для устройств Apple, onnx2tf для TensorFlow, pnnx для NCNN и так далее. У каждого инструмента свой API, свои странности с зависимостями и соглашения об именовании выходных файлов.

Ultralytics поставляет автономные утилиты для экспорта, которые объединяют несколько бэкендов в одном удобном интерфейсе. Ты можешь экспортировать любые torch.nn.Module, включая timm модели изображений, torchvision классификаторы и детекторы, или твои собственные архитектуры, в ONNX, TorchScript, OpenVINO, CoreML, NCNN, PaddlePaddle, MNN, ExecuTorch, и TensorFlow SavedModel без необходимости изучать каждый бэкенд по отдельности.

Зачем использовать Ultralytics для экспорта не-YOLO моделей?

• Единый API для 10 форматов: изучи один стандарт вызова вместо десятка.
• Общий интерфейс утилит: вспомогательные функции для экспорта находятся в ultralytics.utils.export, так что как только пакеты бэкендов установлены, ты можешь использовать один и тот же шаблон вызова для всех форматов.
• Тот же путь кода, что и у экспорта YOLO: те же самые вспомогательные функции лежат в основе каждого экспорта Ultralytics YOLO.
• Квантование FP16 и INT8 встроено для форматов, которые его поддерживают (OpenVINO, CoreML, MNN, NCNN).
• Работает на CPU: для самого шага экспорта GPU не требуется, поэтому ты можешь запустить его локально на любом ноутбуке.

Быстрый старт

Самый быстрый путь — это экспорт в две строки в ONNX без кода YOLO и без настройки, кроме 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")

Поддерживаемые форматы экспорта

Модель torch2* функции принимают стандартный torch.nn.Module и пример входного тензора. MNN, TF SavedModel и TF Frozen Graph проходят через промежуточный артефакт ONNX или Keras. В обоих случаях не требуется никаких специфичных для YOLO атрибутов.

Пошаговые примеры

Каждый пример ниже использует одну и ту же настройку — предобученную ResNet-18 из timm в режиме оценки (evaluation mode):

import timm
import torch

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

Экспорт в ONNX

from ultralytics.utils.export import torch2onnx

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

Для динамического размера батча передай dynamic словарь:

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

По умолчанию opset равен 14, а имя входного тензора по умолчанию "images". Переопредели их с помощью аргументов opset, input_names, или output_names.

Экспорт в TorchScript

Дополнительные зависимости не нужны. Использует torch.jit.trace под капотом.

from ultralytics.utils.export import torch2torchscript

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

Экспорт в OpenVINO

from ultralytics.utils.export import torch2openvino

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

Директория содержит пару файлов с фиксированными именами model.xml и model.bin:

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

Передай dynamic=True для динамических входных форм, half=True для FP16 или int8=True для квантования INT8. Для INT8 дополнительно требуется calibration_dataset.

Требует openvino>=2024.0.0 (или >=2025.2.0 на macOS 15.4+) и torch>=2.1.

Экспорт в 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")

Для classification моделей, передай список имен классов в classifier_names, чтобы добавить классификационную голову к модели CoreML.

Требует coremltools>=9.0, torch>=1.11, и numpy<=2.3.5. Не поддерживается на Windows.

Экспорт в TensorFlow SavedModel

Экспорт в TF SavedModel проходит через ONNX в качестве промежуточного шага:

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")

Функция возвращает модель Keras, а также генерирует файлы TFLite (.tflite) внутри выходной директории:

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

Требования:

• 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 (установи с помощью --extra-index-url https://pypi.ngc.nvidia.com)
• ai-edge-litert>=1.2.0,<1.4.0 на macOS (ai-edge-litert>=1.2.0 на других платформах)
• onnxslim>=0.1.71
• onnx>=1.12.0,<2.0.0
• protobuf>=5

Экспорт в TensorFlow Frozen Graph

Продолжая экспорт SavedModel выше, сконвертируй возвращенную модель Keras в замороженный .pb граф:

from pathlib import Path

from ultralytics.utils.export import keras2pb

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

Экспорт в NCNN

from ultralytics.utils.export import torch2ncnn

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

Директория содержит param и bin файлы с фиксированными именами вместе с Python-оберткой:

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

torch2ncnn() проверяет наличие ncnn и pnnx при первом использовании.

Экспорт в MNN

Экспорт в MNN требует файл ONNX в качестве входного. Сначала экспортируй в ONNX, затем сконвертируй:

from ultralytics.utils.export import onnx2mnn, torch2onnx

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

Поддерживает half=True для FP16 и int8=True для квантования INT8. Требует MNN>=2.9.6 и torch>=1.10.

Экспорт в PaddlePaddle

from ultralytics.utils.export import torch2paddle

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

Директория содержит модель PaddlePaddle и файлы параметров:

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

Требует x2paddle и правильный дистрибутив PaddlePaddle для твоей платформы:

• paddlepaddle-gpu>=3.0.0,<3.3.0 на CUDA
• paddlepaddle==3.0.0 на ARM64 CPU
• paddlepaddle>=3.0.0,<3.3.0 на других CPU

Не поддерживается на NVIDIA Jetson.

Экспорт в ExecuTorch

from ultralytics.utils.export import torch2executorch

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

Экспортируемый .pte файл сохраняется внутри выходной директории:

resnet18_executorch_model/
└── model.pte

Требует torch>=2.9.0 и соответствующая среда выполнения ExecuTorch (pip install executorch). Информацию об использовании среды выполнения см. в интеграции ExecuTorch.

Проверка экспортированной модели

После экспорта проверь числовое соответствие с оригинальной моделью PyTorch перед использованием. Быстрый экспресс-тест с ONNXBackend из ultralytics.nn.backends сравнивает выходные данные и позволяет на ранних этапах выявить ошибки трассировки или квантования:

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

Для других сред выполнения имя входного тензора может отличаться. Например, OpenVINO использует имя аргумента метода forward модели (обычно x для общих моделей), в то время как torch2onnx по умолчанию использует "images".

Известные ограничения

• Поддержка нескольких входов ограничена: torch2onnx и torch2openvino принимает кортеж или список тензоров-примеров для моделей с несколькими входами. torch2torchscript, torch2coreml, torch2ncnn, torch2paddle, и torch2executorch предполагает наличие только одного входного тензора.
• ExecuTorch требует flatc: Для среды выполнения ExecuTorch требуется компилятор FlatBuffers. Установи его с помощью brew install flatbuffers на macOS или apt install flatbuffers-compiler на Ubuntu.
• Отсутствие логического вывода (inference) через Ultralytics: Экспортированные модели, не являющиеся YOLO, нельзя загрузить обратно через YOLO() для выполнения логического вывода. Используй собственную среду выполнения для каждого формата (ONNX Runtime, OpenVINO Runtime, и т. д.).
• Форматы только для YOLO: Axelera и Sony IMX500 экспорты требуют атрибутов модели, специфичных для YOLO, и недоступны для общих моделей.
• Форматы, зависящие от платформы: TensorRT требует GPU NVIDIA. RKNN требует rknn-toolkit2 SDK (только Linux). Edge TPU требует edgetpu_compiler бинарный файл (только Linux).

FAQ

Какие модели я могу экспортировать с помощью Ultralytics?

Любую torch.nn.Module. Сюда входят модели из timm, torchvision или любая пользовательская модель PyTorch. Перед экспортом модель должна быть переведена в режим оценки (model.eval()). ONNX и OpenVINO дополнительно принимают кортеж тензоров-примеров для моделей с несколькими входами.

Какие форматы экспорта работают без GPU?

Все поддерживаемые форматы (TorchScript, ONNX, OpenVINO, CoreML, TF SavedModel, TF Frozen Graph, NCNN, PaddlePaddle, MNN, ExecuTorch) можно экспортировать на CPU. Для самого процесса экспорта GPU не требуется. TensorRT — единственный формат, для которого необходим GPU NVIDIA.

Какая версия Ultralytics мне нужна?

Используй Ultralytics >=8.4.38, который включает модуль ultralytics.utils.export и стандартизированный output_file/output_dir.

Могу ли я экспортировать модель torchvision в CoreML для развертывания на iOS?

Да. Классификаторы, детекторы и модели сегментации torchvision экспортируются в .mlpackage через torch2coreml. Для моделей классификации изображений передай список имен классов в classifier_names, чтобы встроить классификационную головку. Запускай экспорт на macOS или Linux. CoreML не поддерживается в Windows. См. интеграцию CoreML для получения подробной информации о развертывании на iOS.

Могу ли я квантовать мою экспортированную модель до INT8 или FP16?

Да, для нескольких форматов. Передай half=True для FP16 или int8=True для INT8 при экспорте в OpenVINO, CoreML, MNN или NCNN. INT8 в OpenVINO дополнительно требует аргумент calibration_dataset для post-training quantization. См. страницу интеграции каждого формата для уточнения компромиссов при квантовании.