Exportación de PyTorch TorchScript / ONNX

Una vez completado el entrenamiento del modelo, es necesario implementar el modelo en un entorno de producción. PyTorch ofrece múltiples métodos de exportación de modelos, entre los cualesTorchScriptyONNXson los dos formatos más utilizados.

Esta sección detalla los principios, métodos de uso y mejores prácticas de estos dos métodos de exportación.

La exportación de modelos es el proceso de convertir un modelo de PyTorch a un formato que pueda ejecutarse en diferentes plataformas y frameworks. Esto es crucial para escenarios como la implementación de modelos, la inferencia en dispositivos móviles y la migración entre frameworks.


1. Fundamentos de TorchScript

1.1 Qué es TorchScript

TorchScript es un formato de serialización de PyTorch que puede convertir código Python en código de máquina virtual C++ que se puede ejecutar de forma independiente. Los programas TorchScript pueden ejecutarse en entornos sin intérprete de Python.

Las principales características de TorchScript:

  • Convierte gráficos dinámicos en gráficos estáticos.
  • Admite la sintaxis de un subconjunto de Python.
  • Puede ejecutarse en entornos C++.
  • Conserva la estructura y los parámetros del modelo.

1.2 Dos métodos de conversión

TorchScript ofrece dos formas de convertir un modelo a TorchScript:

  • Rastreo de TorchScript (Tracing):genera un gráfico de cálculo estático al ejecutar el modelo y registrar las operaciones.
  • Scripting de TorchScript (Scripting):analiza directamente el código Python y lo compila a TorchScript.

2. Rastreo de TorchScript (Tracing)

2.1 Métodos básicos de rastreo

Ejemplo

import torch
import torch.nn as nn

# ── Definir el modelo ──────────────────────────────────────
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(128 * 8 * 8, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x


# Crear instancia del modelo
model = SimpleNet()
model.eval()

# Entrada de ejemplo
example_input = torch.randn(1, 3, 64, 64)

# ── Método 1: rastreo de torch.jit.trace ───────────────────
# Crear TorchScript ejecutando el modelo y registrando las operaciones
traced_model = torch.jit.trace(model, example_input)

print("Modelo rastreado:")
print(traced_model)

# Guardar modelo
traced_model.save("simple_net_traced.pt")

# Cargar modelo
loaded_model = torch.jit.load("simple_net_traced.pt")

# Usar el modelo cargado para inferencia
output = loaded_model(example_input)
print(f"Forma de salida: {output.shape}")

2.2 Limitaciones del rastreo

El método de rastreo tiene algunas limitaciones:

  • Solo registra las operaciones realmente ejecutadas.
  • El flujo de control (como if, for) queda fijo.
  • Las entradas de tamaño dinámico pueden causar problemas.

Ejemplo

# Ejemplo de limitaciones del rastreo

class DynamicModel(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        # El flujo de control se fija durante el seguimiento
        if x.sum() > 0:
            return x * 2
        else:
            return x / 2


model = DynamicModel()
model.eval()

# Al rastrear, la rama if se fija a la ruta del seguimiento
example_input = torch.tensor([1.0])
traced = torch.jit.trace(model, example_input)

# Incluso si la entrada es diferente, se ejecuta la misma rama
print(traced(torch.tensor([1.0])))  # 2
print(traced(torch.tensor([-1.0]))) # Sigue siendo 2, no -0.5

Para modelos con flujo de control, se debe usar Scripting de TorchScript en lugar de Tracing.


3. Scripting de TorchScript

3.1 Uso básico

Ejemplo

import torch

# ── Usando el decorador @torch.jit.script ───────────────
@torch.jit.script
def scripted_function(x: torch.Tensor) -> torch.Tensor:
    """Usar el método de script para convertir la función"""
    if x.sum() > 0:
        return x * 2
    else:
        return x / 2


# Probar funciones de script
input1 = torch.tensor([1.0, 2.0])
input2 = torch.tensor([-1.0, -2.0])

print(scripted_function(input1))  # [2., 4.]
print(scripted_function(input2))    # [-0.5, -1.]


# ── Modelo de script ───────────────────────────────────
class ScriptableModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 10)

    @torch.jit.export
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.relu(self.fc(x))

    @torch.jit.export
    def predict(self, x: torch.Tensor) -> torch.Tensor:
        """Métodos de exportación adicionales"""
        out = self.forward(x)
        return torch.argmax(out, dim=1)


model = ScriptableModel()

# Modelo scripteado
scripted_model = torch.jit.script(model)
print(scripted_model)

# Guardar
scripted_model.save("scripted_model.pt")

3.2 Scripting de modelos complejos

Ejemplo

# Un ejemplo de script más complejo: modelo con condiciones

class ConditionModel(nn.Module):
    def __init__(self, num_classes: int):
        super().__init__()
        self.num_classes = num_classes
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten()
        )
        self.classifier = nn.Linear(64, num_classes)

    def forward(self, x: torch.Tensor, use_softmax: bool = False) -> torch.Tensor:
        """
Modelos que admiten condiciones dinámicas
        """

        features = self.features(x)
        logits = self.classifier(features)

        if use_softmax:
            return torch.softmax(logits, dim=1)
        else:
            return logits

    def get_prediction(self, x: torch.Tensor) -> torch.Tensor:
        """Métodos auxiliares"""
        logits = self.forward(x, use_softmax=False)
        return torch.argmax(logits, dim=1)


# Scripting
model = ConditionModel(num_classes=10)
scripted_model = torch.jit.script(model, example_inputs=(torch.randn(1, 3, 32, 32),))

# Prueba
test_input = torch.randn(2, 3, 32, 32)
output1 = scripted_model(test_input, use_softmax=False)
output2 = scripted_model(test_input, use_softmax=True)

print(f"Forma de salida de los logits: {output1.shape}")
print(f"Forma de salida de Softmax: {output2.shape}")

4. Exportación a ONNX

4.1 Fundamentos de ONNX

ONNX(Open Neural Network eXchange)ONNX es un formato abierto de intercambio de redes neuronales que admite la conversión de modelos entre diferentes frameworks de aprendizaje profundo.

Las ventajas de ONNX:

  • Multi-framework: compatible con PyTorch, TensorFlow, Caffe2, etc.
  • Multiplataforma: admite CPU, GPU, móvil y otras plataformas.
  • Optimización de hardware: se puede utilizar ONNX Runtime para una inferencia eficiente.
  • Ricas herramientas: hay una gran cantidad de herramientas de optimización y soluciones de implementación.

4.2 Exportación básica a ONNX

Ejemplo

import torch
import torch.nn as nn
import torchvision

# ── Definir el modelo ──────────────────────────────────────
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten()
        )
        self.classifier = nn.Linear(64, 10)

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x


model = ImageClassifier()
model.eval()

# Entrada de ejemplo
example_input = torch.randn(1, 3, 32, 32)

# ── Exportar a ONNX ──────────────────────────────────
output_path = "image_classifier.onnx"

torch.onnx.export(
    model,
    example_input,
    output_path,
    export_params=True,        # Exportar parámetros del modelo
    opset_version=14,          # Versión de ONNX
    do_constant_folding=True, # Optimización de plegado de constantes
    input_names=['input'],     # Nombre del tensor de entrada
    output_names=['output'],   # Nombre del tensor de salida
    dynamic_axes={
        'input': {0: 'batch_size'},    # Dimensión de lote dinámica
        'output': {0: 'batch_size'}
    }
)

print(f"Modelo exportado a: {output_path}")

# Verificar el modelo exportado
import onnx
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print(¡Modelo ONNX validado correctamente!)

4.3 Exportar modelos complejos

Ejemplo

# Exportar el modelo ResNet completo

import torchvision.models as models

# Cargar el modelo preentrenado
model = models.resnet18(pretrained=True)
model.eval()

# Ejemplo de entrada (224x224 estándar de ResNet)
example_input = torch.randn(1, 3, 224, 224)

# Exportar
torch.onnx.export(
    model,
    example_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=14,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size'}
    }
)

print(ResNet18 exportado)

# Manejar capas especiales al exportar
# Para las capas que requieren un manejo especial, usa workaround
class ModelWithSpecialOps(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3, padding=1)

    def forward(self, x):
        # Usa F.interpolate en lugar del alias de nn.functional
        x = self.conv(x)
        x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
        return x


model = ModelWithSpecialOps()
model.eval()
torch.onnx.export(
    model,
    torch.randn(1, 3, 32, 32),
    "special_ops.onnx",
    input_names=['input'],
    output_names=['output'],
    opset_version=14
)

5. Validación y optimización de la exportación

5.1 Validación del modelo exportado

Ejemplo

import numpy as np
import onnx
import onnxruntime as ort

def verify_onnx_model(onnx_path, pytorch_model, test_input):
    """
Verificar la consistencia de las salidas entre el modelo ONNX y el modelo PyTorch.
    """

    # 1. Verificar el modelo ONNX
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)
    print("✓ Validación de la estructura del modelo ONNX superada")

    # 2. Inferencia con PyTorch
    pytorch_model.eval()
    with torch.no_grad():
        pytorch_output = pytorch_model(test_input)

    # 3. Inferencia con ONNX Runtime
    ort_session = ort.InferenceSession(onnx_path)
    ort_output = ort_session.run(None, {'input': test_input.numpy()})[0]

    # 4. Comparar salidas
    diff = np.abs(pytorch_output.numpy() - ort_output).max()
    print(f"✓ Diferencia máxima de salida: {diff:.6f}")

    if diff < 1e-5:
        print("✓ Verificación de salida aprobada")
        return True
    else:
        print("&#x26a0; La diferencia de salida es grande")
        return False


# Ejemplo de uso
model = ImageClassifier()
model.eval()

# Primero exportar
torch.onnx.export(
    model,
    torch.randn(1, 3, 32, 32),
    "test_model.onnx",
    input_names=['input'],
    output_names=['output']
)

# Verificar
test_input = torch.randn(2, 3, 32, 32)
verify_onnx_model("test_model.onnx", model, test_input)

5.2 Optimización con ONNX Runtime

Ejemplo

import onnxruntime as ort
from onnxruntime import GraphOptimizationLevel

# Crear sesión de inferencia y configurar optimización
def create_optimized_session(onnx_path, providers=None):
    if providers is None:
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']

    # Configurar opciones de sesión
    sess_options = ort.SessionOptions()

    # Habilitar optimización de grafos
    sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL

    # Habilitar otras optimizaciones
    sess_options.intra_op_num_threads = 4
    sess_options.inter_op_num_threads = 4

    # Crear sesión
    session = ort.InferenceSession(onnx_path, sess_options, providers=providers)

    return session


# Usar la sesión optimizada para inferencia
session = create_optimized_session("resnet18.onnx")

# Obtener nombres de entrada/salida
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Inferencia
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = session.run([output_name], {input_name: input_data})[0]

print(f"Forma de salida de inferencia: {output.shape}")

5.3 Cuantización de modelos

Ejemplo

# Cuantización de modelos ONNX

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

def quantize_onnx_model(input_path, output_path):
    """
Cuantización dinámica de modelos ONNX
Reducir el tamaño del modelo y acelerar la inferencia
    """

    # Cuantización dinámica (no requiere datos de calibración)
    quantize_dynamic(
        input_path,
        output_path,
        # Especificar los tipos de pesos a cuantizar
        weight_type=QuantType.QInt8
    )

    # Obtener comparación de tamaños de archivo
    import os
    original_size = os.path.getsize(input_path) / 1024 / 1024
    quantized_size = os.path.getsize(output_path) / 1024 / 1024

    print(f"Modelo original: {original_size:.2f} MB")
    print(f"Modelo cuantizado: {quantized_size:.2f} MB")
    print(f"Relación de compresión: {original_size/quantized_size:.2f}x")


# Uso
quantize_onnx_model("resnet18.onnx", "resnet18_quantized.onnx")

6. Implementación en dispositivos móviles

6.1 Exportación a formatos móviles

Ejemplo

# Método 1: Implementación móvil con TorchScript
# 1. Rastrear el modelo
model = ImageClassifier()
model.eval()
example_input = torch.randn(1, 3, 32, 32)

traced_model = torch.jit.trace(model, example_input)

# 2. Optimizar el modelo (optimización móvil)
optimized_model = torch.jit.optimize_for_inference(traced_model)

# 3. Guardar el modelo móvil
optimized_model.save("mobile_model.pt")

# Método 2: Usar torchmobile (si está disponible)
# torchmobile se usa para implementación móvil más ligera
# Consulte la documentación oficial para compilación cruzada

# Método 3: Exportar a ONNX y usar runtime móvil
# iOS: Usar Core ML Tools
# Android: Usar ONNX Runtime Mobile

print("Modelo móvil listo")

6.2 Exportación a Core ML (iOS)

Para la plataforma iOS, puede usar Core ML Tools para convertir el modelo al formato Core ML. Primero debe convertirlo al formato ONNX y luego convertirlo con Core ML Tools.

Ejemplo

# Requiere instalar coremltools
# pip install coremltools

# Pasos para exportar a Core ML
# Nota: esto solo está disponible en macOS

# Paso 1: Primero exportar el modelo PyTorch a ONNX
# from pytorch_example import your_model
# model = your_model()
# model.eval()
# example_input = torch.randn(1, 3, 224, 224)
# torch.onnx.export(
#     model,
#     example_input,
#     "model.onnx",
#     input_names=['input'],
#     output_names=['output']
# )

# Paso 2: Usar Core ML Tools para convertir al formato Core ML
# from coremltools.converters import onnx as onnx_coreml
# coreml_model = onnx_coreml.convert(
#     model="model.onnx",
#     minimum_deployment_target='13',
#     image_input_names=['input']
# )
# coreml_model.save("model.mlmodel")

# Ejemplo de uso
print("La implementación en iOS requiere entorno macOS")
print("Pasos detallados: 1. pip install coremltools 2. Exportar ONNX 3. Convertir con coremltools")

7. Mejores prácticas y problemas comunes

7.1 Solución de problemas de exportación

Problema Causa Solución
Error de exportación Operaciones no compatibles Usar una versión más alta de opset_version o reemplazar las operaciones
Salida inconsistente Flujo de control dinámico Usar scripting en lugar de tracing, o fijar el tamaño de entrada
Error de ONNX Runtime Operadores no compatibles Consultar la lista de operaciones compatibles con ONNX

7.2 Comparación de formatos de exportación

Formato Ventajas Desventajas Escenarios aplicables
TorchScript (.pt) Compatibilidad nativa con PyTorch Escasa compatibilidad entre frameworks Implementación en escritorio/servidor
ONNX (.onnx) Multi-framework, multiplataforma Algunas operaciones no son compatibles Implementación general

7.3 Lista de verificación de exportación

Ejemplo

# Lista de verificación antes de exportar

def export_checklist(model, example_input):
    """Verificación antes de exportar"""
    model.eval()

    # 1. Asegurarse de que el modelo esté en modo de inferencia
    print(f"Modo del modelo: {'eval' if not model.training else 'train'}")

    # 2. Verificar el tamaño de entrada
    print(f"Forma de la entrada de ejemplo: {example_input.shape}")

    # 3. Verificar la propagación hacia adelante
    with torch.no_grad():
        output = model(example_input)
    print(f"Forma de salida: {output.shape}")

    # 4. Verificar si hay operaciones no serializables
    # Por ejemplo: funciones lambda
    for name, module in model.named_modules():
        if hasattr(module, '__call__'):
            pass  # Verificar módulos personalizados

    print("✓ Verificación de exportación completada")


# Ejemplo de uso
model = ImageClassifier()
example_input = torch.randn(1, 3, 32, 32)
export_checklist(model, example_input)

7.4 Flujo de exportación completo

Ejemplo

# Ejemplo de flujo completo de exportación de modelos

class ProductionModel(nn.Module):
    """Modelo de nivel de producción"""
    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(3, stride=2, padding=1),
            # ... más capas
        )
        self.head = nn.Linear(512, 10)

    def forward(self, x, return_features=False):
        features = self.backbone(x)
        features = torch.nn.functional.adaptive_avg_pool2d(features, 1).flatten(1)

        if return_features:
            return features

        return self.head(features)


# Flujo de exportación completo
model = ProductionModel()
model.load_state_dict(torch.load("weights.pth"))  # Cargar pesos
model.eval()

# Preparar entrada de ejemplo
example_input = torch.randn(1, 3, 224, 224)

# 1. Hacer tracing de TorchScript
traced = torch.jit.trace(model, example_input)
traced.save("model_traced.pt")

# 2. Exportar ONNX
torch.onnx.export(
    traced,
    example_input,
    "model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
    opset_version=14
)

# 3. Optimizar ONNX
# Usar onnx-simplifier para eliminar operaciones redundantes
# onnxsim model.onnx model_simplified.onnx

print("¡Exportación de todos los formatos completada!")
Otras extensiones