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.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
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
# ── 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
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.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
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 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("⚠ 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
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
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
# 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
# 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
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
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!")