Esportare modelli con TorchScript

Modelli di AI scalabili con PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Cos'è TorchScript?

Indipendente da Python
Efficiente in produzione
Esempi:
- Deploy su dispositivi mobili

Un telefono cellulare

Cos'è TorchScript?

Indipendente da Python
Efficiente in produzione
Esempi:
- Deploy su dispositivi mobili
- Inferenza ad alte prestazioni in produzione

Un telefono e un laptop con ingranaggi che rappresentano un ambiente di produzione

Convertire modelli in TorchScript

Due metodi di conversione:

torch.jit.trace: usa input di esempio per tracciare l'esecuzione
torch.jit.script: compila analizzando il codice Python

Quando usarli:

Usa trace per modelli semplici
Usa script per modelli con control flow (es. loop)

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)

Salvare e caricare modelli TorchScript

Salvataggio del modello:
- torch.jit.save: salva il modello scriptato su file
Caricamento del modello:
- torch.jit.load: ricarica il modello per l'inferenza

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

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

Fare inferenza con TorchScript

Passi:
- Carica il modello TorchScript
- Passa gli input al modello per le predizioni
- Output identici a PyTorch

Input di esempio:

Tensore input: [1.0, 2.0, 3.0]

Output di esempio:

Tensore output: [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 in breve

torch.jit.trace: per modelli statici
torch.jit.script: gestisce il control flow dinamico
torch.jit.save: salva il modello scriptato
torch.jit.load: ricarica per l'inferenza

Ayo berlatih!

Modelli di AI scalabili con PyTorch Lightning