Exportar modelos con TorchScript

Modelos de IA escalables con PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

¿Qué es TorchScript?

Independiente de Python
Eficiente en producción
Ejemplos:
- Despliegue en móviles

Un teléfono móvil

¿Qué es TorchScript?

Independiente de Python
Eficiente en producción
Ejemplos:
- Despliegue en móviles
- Inferencia de alto rendimiento en producción

Un teléfono móvil y un portátil con engranajes que representan un entorno de producción

Convertir modelos a TorchScript

Dos métodos de conversión:

torch.jit.trace: usa entradas de ejemplo para trazar la ejecución
torch.jit.script: compila analizando el código Python

Cuándo usar:

Usa trace para modelos simples
Usa script para modelos con control de flujo (p. ej., bucles)

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

model = SimpleModel()
scripted_model = torch.jit.script(model)

Guardar y cargar modelos TorchScript

Guardar el modelo:
- torch.jit.save: guarda el modelo con script en un archivo
Cargar el modelo:
- torch.jit.load: vuelve a cargar el modelo para inferencia

# Save the model
torch.jit.save(scripted_mod,"model.pt")

# Load the model
loaded_model=torch.jit.load("model.pt")

Inferencia con TorchScript

Pasos:
- Carga el modelo TorchScript
- Pasa entradas al modelo para predecir
- Salidas idénticas a PyTorch

Entrada de ejemplo:

Tensor de entrada: [1.0, 2.0, 3.0]

Salida de ejemplo:

Tensor de salida: [2.0, 4.0, 6.0]

# Perform inference
input_arr = [1.0, 2.0, 3.0]
input_tensor = torch.tensor(input_arr)

output = loaded_model(input_tensor)
print(output)

TorchScript en pocas palabras

torch.jit.trace: para modelos estáticos
torch.jit.script: maneja control de flujo dinámico
torch.jit.save: guarda el modelo con script
torch.jit.load: recarga para inferencia

¡Vamos a practicar!

Modelos de IA escalables con PyTorch Lightning