Export modelů pomocí TorchScript

Škálovatelné modely AI s PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Co je TorchScript?

Nezávislý na Pythonu
Efektivní v produkci
Příklady:
- Nasazení na mobilních zařízeních

Mobilní telefon

Co je TorchScript?

Nezávislý na Pythonu
Efektivní v produkci
Příklady:
- Nasazení na mobilních zařízeních
- Vysoce výkonná inference v produkčních systémech

Mobilní telefon a laptop s ozubenými koly představující produkční prostředí

Převod modelů na TorchScript

Dvě metody převodu:

torch.jit.trace: Sleduje provádění na základě ukázkových vstupů
torch.jit.script: Kompiluje model analýzou kódu Python

Kdy použít:

trace pro jednodušší modely
script pro modely s řídicím tokem (např. smyčky)

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)

Ukládání a načítání modelů TorchScript

Uložení modelu:
- torch.jit.save: Uloží skriptovaný model do souboru
Načtení modelu:
- torch.jit.load: Načte model zpět pro inferenci

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

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

Inference s TorchScript

Postup:
- Načtěte model TorchScript
- Předejte vstupy modelu pro predikce
- Výstupy jsou totožné s výstupy PyTorch

Příklad vstupu:

Vstupní tensor: [1.0, 2.0, 3.0]

Příklad výstupu:

Výstupní 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 stručně

torch.jit.trace: Vhodné pro statické modely
torch.jit.script: Zvládá dynamický řídicí tok
torch.jit.save: Uloží skriptovaný model
torch.jit.load: Načte model pro inferenci

Pojďme procvičovat!

Škálovatelné modely AI s PyTorch Lightning