Как экспортироватьPyTorch YOLO , с помощью Ultralytics

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

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

Почему стоит использовать Ultralytics YOLO ?

• Один API для 10 форматов: освойте один способ вызова вместо десятка.
• Общая поверхность для коммуникаций: помощники по экспорту находятся в ultralytics.utils.export, поэтому после установки пакетов бэкэнда вы сможете использовать одинаковую схему вызова для всех форматов.
• Тот же алгоритм, что и YOLO : все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 и пример tensor входных данных. MNN, TF SavedModel и TF Graph проходят через промежуточный артефакт ONNX Keras. В обоих случаях не требуется никаких атрибутов, YOLO.

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

Во всех приведенных ниже примерах используется одна и та же конфигурация: предварительно обученная модель ResNet-18 из пакета timm в режиме оценки:

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

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

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

Экспорт в форматSavedModel TensorFlow

SavedModel TF SavedModel ONNX промежуточного этапа используется формат 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 граф TensorFlow

В продолжение описанного выше 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 экспортируйте в 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 и параметров PaddlePaddle :

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

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

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

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

Экспорт в 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.forward(im)[0]

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

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

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

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

Часто задаваемые вопросы

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

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

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

Все поддерживаемые форматы (TorchScript, ONNX, OpenVINO, CoreML, TF SavedModel, TF Graph, NCNN, PaddlePaddle, MNN, ExecuTorch) можно экспортировать на CPU. Для самого процесса экспорта GPU не 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 не CoreML в Windows. См. CoreML подробности о iOS на iOS .

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

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

📅 Создано 0 дней назад ✏️ Обновлено 0 дней назад

Комментарии