Como exportar modelos PyTorch que não são YOLO com Ultralytics

Implementar modelos PyTorch em produção geralmente significa lidar com um exportador diferente para cada destino: torch.onnx.export para ONNX, coremltools para dispositivos Apple, onnx2tf para TensorFlow, pnnx para NCNN, e assim por diante. Cada ferramenta tem sua própria API, peculiaridades de dependência e convenções de saída.

O Ultralytics oferece utilitários de exportação independentes que envolvem vários backends sob uma interface consistente. Você pode exportar qualquer modelo de imagem torch.nn.ModuleUma limitação significativa de muitos modelos de detecção especializados é o seu foco restrito. Dentro do ecossistema Ultralytics, não estás limitado apenas à detecção de objetos. As ferramentas estendem-se perfeitamente a múltiplas timm, classificadores e detectores do torchvision, ou suas próprias arquiteturas personalizadas para ONNX, TorchScript, OpenVINO, CoreML, NCNN, PaddlePaddle, MNN, ExecuTorch, e TensorFlow SavedModel sem precisar aprender cada backend separadamente.

Por que usar o Ultralytics para exportação de modelos não YOLO?

• Uma única API para 10 formatos: aprenda uma única convenção de chamada em vez de uma dúzia.
• Interface utilitária compartilhada: os auxiliares de exportação ficam em ultralytics.utils.export, então, uma vez que os pacotes de backend estejam instalados, você pode manter o mesmo padrão de chamada entre formatos.
• O mesmo caminho de código das exportações YOLO: os mesmos auxiliares alimentam toda exportação YOLO do Ultralytics.
• Quantização FP16 e INT8 integradas para formatos que as suportam (OpenVINO, CoreML, MNN, NCNN).
• Funciona na CPU: não é necessária GPU para a etapa de exportação em si, então você pode executá-la localmente em qualquer laptop.

Início Rápido

O caminho mais rápido é uma exportação de duas linhas para ONNX sem código YOLO e sem configuração além do 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")

Formatos de Exportação Suportados

A classe de modelo torch2* As funções aceitam um torch.nn.Module padrão e um tensor de entrada de exemplo. MNN, TF SavedModel e TF Frozen Graph passam por um artefato intermediário ONNX ou Keras. Nenhum atributo específico do YOLO é necessário em ambos os casos.

Exemplos passo a passo

Cada exemplo abaixo usa a mesma configuração: um ResNet-18 pré-treinado do timm em modo de avaliação:

import timm
import torch

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

Exportar para ONNX

from ultralytics.utils.export import torch2onnx

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

Para tamanho de lote (batch size) dinâmico, passe um dicionário dynamic:

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

O opset padrão é 14 e o nome de entrada padrão é "images". Substitua com os argumentos opset, input_names, ou output_names.

Exportar para TorchScript

Nenhuma dependência extra necessária. Usa torch.jit.trace por baixo dos panos.

from ultralytics.utils.export import torch2torchscript

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

Exportar para OpenVINO

from ultralytics.utils.export import torch2openvino

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

O diretório contém um par de model.xml e model.bin de nome fixo:

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

Passe dynamic=True para formas de entrada dinâmicas, half=True para FP16, ou int8=True para quantização INT8. INT8 requer adicionalmente um calibration_dataset.

Requer openvino>=2024.0.0 (ou >=2025.2.0 no macOS 15.4+) e torch>=2.1.

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

Para classification modelos, passe uma lista de nomes de classe para classifier_names para adicionar um cabeçalho de classificação ao modelo CoreML.

Requer coremltools>=9.0, torch>=1.11, e numpy<=2.3.5. Não é suportado no Windows.

Exportar para TensorFlow SavedModel

A exportação para TF SavedModel passa pelo ONNX como um passo intermediário:

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

A função retorna um modelo Keras e também gera arquivos TFLite (.tflite) dentro do diretório de saída:

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

Requisitos:

• 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 (instale com --extra-index-url https://pypi.ngc.nvidia.com)
• ai-edge-litert>=1.2.0,<1.4.0 no macOS (ai-edge-litert>=1.2.0 em outras plataformas)
• onnxslim>=0.1.71
• onnx>=1.12.0,<2.0.0
• protobuf>=5

Exportar para TensorFlow Frozen Graph

Continuando a partir da exportação do SavedModel acima, converta o modelo Keras retornado para um grafo .pb congelado:

from pathlib import Path

from ultralytics.utils.export import keras2pb

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

Exportar para NCNN

from ultralytics.utils.export import torch2ncnn

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

O diretório contém arquivos param e bin de nome fixo, juntamente com um wrapper em Python:

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

torch2ncnn() verifica por ncnn e pnnx do Ultralytics no primeiro uso.

Exportar para MNN

A exportação para MNN requer um arquivo ONNX como entrada. Exporte para ONNX primeiro e, em seguida, converta:

from ultralytics.utils.export import onnx2mnn, torch2onnx

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

Suporta half=True para FP16 e int8=True para quantização INT8. Requer MNN>=2.9.6 e torch>=1.10.

Exportar para PaddlePaddle

from ultralytics.utils.export import torch2paddle

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

O diretório contém o modelo PaddlePaddle e os arquivos de parâmetro:

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

Requer x2paddle e a distribuição PaddlePaddle correta para a sua plataforma:

• paddlepaddle-gpu>=3.0.0,<3.3.0 em CUDA
• paddlepaddle==3.0.0 em CPU ARM64
• paddlepaddle>=3.0.0,<3.3.0 em outras CPUs

Não é suportado no NVIDIA Jetson.

Exportar para ExecuTorch

from ultralytics.utils.export import torch2executorch

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

O arquivo .pte exportado é salvo dentro do diretório de saída:

resnet18_executorch_model/
└── model.pte

Requer torch>=2.9.0 e um runtime ExecuTorch correspondente (pip install executorch). Para uso do runtime, veja a integração com ExecuTorch.

Verifique o seu modelo exportado

Após a exportação, verifique a paridade numérica com o modelo PyTorch original antes de fazer o envio. Um rápido teste de sanidade com ONNXBackend de ultralytics.nn.backends compara as saídas e sinaliza erros de rastreamento ou quantização precocemente:

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

Para outros runtimes, o nome do tensor de entrada pode diferir. O OpenVINO, por exemplo, usa o nome do argumento forward do modelo (tipicamente x para modelos genéricos), enquanto o torch2onnx usa por padrão "images".

Limitações conhecidas

• Suporte para múltiplas entradas é irregular: torch2onnx e torch2openvino aceitam uma tupla ou lista de tensores de exemplo para modelos com múltiplas entradas. torch2torchscript, torch2coreml, torch2ncnn, torch2paddle, e torch2executorch assumem um único tensor de entrada.
• O ExecuTorch precisa do flatc: O runtime ExecuTorch requer o compilador FlatBuffers. Instale com brew install flatbuffers no macOS ou apt install flatbuffers-compiler no Ubuntu.
• Sem inferência via Ultralytics: Modelos não-YOLO exportados não podem ser carregados de volta através do YOLO() para inferência. Use o runtime nativo de cada formato (ONNX Runtime, OpenVINO Runtime, etc.).
• Formatos exclusivos YOLO: Axelera e Sony IMX500 as exportações requerem atributos de modelo específicos do YOLO e não estão disponíveis para modelos genéricos.
• Formatos específicos de plataforma: TensorRT requer uma GPU NVIDIA. RKNN requer o rknn-toolkit2 SDK (apenas Linux). Edge TPU requer o edgetpu_compiler binário (apenas Linux).

FAQ

Quais modelos posso exportar com a Ultralytics?

Qualquer torch.nn.Module. Isso inclui modelos de timm, torchvision ou qualquer modelo PyTorch personalizado. O modelo deve estar em modo de avaliação (model.eval()) antes da exportação. ONNX e OpenVINO aceitam adicionalmente uma tupla de tensores de exemplo para modelos de múltiplas entradas.

Quais formatos de exportação funcionam sem uma GPU?

Todos os formatos suportados (TorchScript, ONNX, OpenVINO, CoreML, TF SavedModel, TF Frozen Graph, NCNN, PaddlePaddle, MNN, ExecuTorch) podem exportar na CPU. Nenhuma GPU é necessária para o processo de exportação em si. TensorRT é o único formato que requer uma GPU NVIDIA.

Qual versão da Ultralytics eu preciso?

Use a Ultralytics >=8.4.38, que inclui o módulo ultralytics.utils.export e a padronizada output_file/output_dir.

Posso exportar um modelo torchvision para CoreML para implantação em iOS?

Sim. Classificadores, detectores e modelos de segmentação da torchvision exportam para .mlpackage via torch2coreml. Para modelos de classificação de imagem, passe uma lista de nomes de classes para classifier_names para integrar um head de classificação. Execute a exportação no macOS ou Linux. CoreML não é suportado no Windows. Veja a integração com CoreML para detalhes de implantação em iOS.

Posso quantizar meu modelo exportado para INT8 ou FP16?

Sim, para vários formatos. Passe half=True para FP16 ou int8=True para INT8 ao exportar para OpenVINO, CoreML, MNN ou NCNN. INT8 no OpenVINO requer adicionalmente um argumento calibration_dataset para quantização pós-treinamento. Veja a página de integração de cada formato para os trade-offs de quantização.