Modelle mit TorchScript exportieren

Skalierbare KI-Modelle mit PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Was ist TorchScript?

Unabhängig von Python
Effizient in Produktion
Beispiele:
- Deployment auf Mobilgeräten

Ein Smartphone

Was ist TorchScript?

Unabhängig von Python
Effizient in Produktion
Beispiele:
- Deployment auf Mobilgeräten
- Hochperformante Inferenz in Produktionssystemen

Ein Smartphone und ein Laptop mit Zahnrädern für eine Produktionsumgebung

Modelle nach TorchScript konvertieren

Zwei Konvertierungswege:

torch.jit.trace: Nutzt Beispiel-Eingaben zum Tracen
torch.jit.script: Kompiliert per Codeanalyse

Wann verwenden:

trace für einfache Modelle
script für Modelle mit Kontrollfluss (z. B. Schleifen)

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)

TorchScript-Modelle speichern und laden

Modell speichern:
- torch.jit.save: Skriptiertes Modell in Datei speichern
Modell laden:
- torch.jit.load: Modell fürs Inferenzieren laden

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

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

Inferenz mit TorchScript durchführen

Schritte:
- TorchScript-Modell laden
- Eingaben für Vorhersagen übergeben
- Ausgaben sind identisch zu PyTorch

Beispiel-Eingabe:

Eingabe-Tensor: [1.0, 2.0, 3.0]

Beispiel-Ausgabe:

Ausgabe-Tensor: [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 auf den Punkt

torch.jit.trace: Für statische Modelle
torch.jit.script: Für dynamische Kontrollflüsse
torch.jit.save: Speichert das skriptierte Modell
torch.jit.load: Lädt für Inferenz

Lass uns üben!

Skalierbare KI-Modelle mit PyTorch Lightning