Exporter des modèles avec TorchScript

Modèles d'IA évolutifs avec PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Qu’est-ce que TorchScript ?

Indépendant de Python
Efficace en production
Exemples :
- Déploiement sur mobile

Un téléphone mobile

Qu’est-ce que TorchScript ?

Indépendant de Python
Efficace en production
Exemples :
- Déploiement sur mobile
- Inférence haute performance en production

Un téléphone mobile et un ordinateur portable avec des engrenages représentant un environnement de production

Convertir des modèles en TorchScript

Deux méthodes de conversion :

torch.jit.trace : suit l’exécution avec des entrées d’exemple
torch.jit.script : compile en analysant le code Python

Quand les utiliser :

Utiliser trace pour les modèles simples
Utiliser script pour les modèles avec contrôle de flux (boucles, etc.)

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)

Sauvegarde et chargement des modèles TorchScript

Sauvegarder le modèle :
- torch.jit.save : enregistre le modèle scripté dans un fichier
Charger le modèle :
- torch.jit.load : recharge le modèle pour l’inférence

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

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

Réaliser l’inférence avec TorchScript

Étapes :
- Charger le modèle TorchScript
- Fournir les entrées pour prédire
- Sorties identiques à PyTorch

Exemple d’entrée :

Tenseur d’entrée : [1.0, 2.0, 3.0]

Exemple de sortie :

Tenseur de sortie : [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 bref

torch.jit.trace : pour modèles statiques
torch.jit.script : gère le contrôle de flux dynamique
torch.jit.save : enregistre le modèle scripté
torch.jit.load : recharge pour l’inférence

Passons à la pratique !

Modèles d'IA évolutifs avec PyTorch Lightning