Eksportowanie modeli za pomocą TorchScript

Skalowalne modele AI z PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Czym jest TorchScript?

Niezależny od Pythona
Wydajny w produkcji
Przykłady:
- Wdrożenie na urządzeniach mobilnych

Telefon komórkowy

Czym jest TorchScript?

Niezależny od Pythona
Wydajny w produkcji
Przykłady:
- Wdrożenie na urządzeniach mobilnych
- Wysokowydajne wnioskowanie w systemach produkcyjnych

Telefon komórkowy i laptop z przekładniami symbolizującymi środowisko produkcyjne

Konwersja modeli do TorchScript

Dwie metody konwersji:

torch.jit.trace: Śledzi wykonanie na podstawie przykładowych danych
torch.jit.script: Kompiluje model analizując kod Pythona

Kiedy używać:

trace – dla prostszych modeli
script – dla modeli z przepływem sterowania (np. pętle)

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)

Zapisywanie i wczytywanie modeli TorchScript

Zapisywanie modelu:
- torch.jit.save: Zapisuje model skryptowy do pliku
Wczytywanie modelu:
- torch.jit.load: Wczytuje model do wnioskowania

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

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

Wnioskowanie z użyciem TorchScript

Kroki:
- Wczytaj model TorchScript
- Przekaż dane wejściowe do modelu
- Wyniki są identyczne jak w PyTorch

Przykładowe dane wejściowe:

Tensor wejściowy: [1.0, 2.0, 3.0]

Przykładowe dane wyjściowe:

Tensor wyjściowy: [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 w skrócie

torch.jit.trace: Działa dla modeli statycznych
torch.jit.script: Obsługuje dynamiczny przepływ sterowania
torch.jit.save: Zapisuje model skryptowy
torch.jit.load: Wczytuje model do wnioskowania

Czas na ćwiczenia!

Skalowalne modele AI z PyTorch Lightning