Exportando modelos com TorchScript

Modelos de IA escaláveis com PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

O que é TorchScript?

Independente do Python
Eficiente em produção
Exemplos:
- Deploy em dispositivos móveis

Um celular

O que é TorchScript?

Independente do Python
Eficiente em produção
Exemplos:
- Deploy em dispositivos móveis
- Inferência de alta performance em produção

Um celular e um laptop com engrenagens representando um ambiente de produção

Convertendo modelos para TorchScript

Duas formas de conversão:

torch.jit.trace: usa entradas de exemplo para rastrear a execução
torch.jit.script: compila analisando o código Python

Quando usar:

Use trace para modelos simples
Use script para modelos com controle de fluxo (p.ex., loops)

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)

Salvando e carregando modelos TorchScript

Salvar o modelo:
- torch.jit.save: salva o modelo scriptado em um arquivo
Carregar o modelo:
- torch.jit.load: carrega o modelo para inferência

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

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

Fazendo inferência com TorchScript

Etapas:
- Carrega o modelo TorchScript
- Passa entradas para obter previsões
- Saídas idênticas às do PyTorch

Entrada de exemplo:

Tensor de entrada: [1.0, 2.0, 3.0]

Saída de exemplo:

Tensor de saída: [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)

Resumo do TorchScript

torch.jit.trace: para modelos estáticos
torch.jit.script: lida com controle de fluxo dinâmico
torch.jit.save: salva o modelo scriptado
torch.jit.load: recarrega para inferência

Vamos praticar!

Modelos de IA escaláveis com PyTorch Lightning