Modellen exporteren met TorchScript

Schaalbare AI-modellen met PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Wat is TorchScript?

Onafhankelijk van Python
Efficiënt in productie
Voorbeelden:
- Uitrol op mobiele apparaten

Een mobiele telefoon

Wat is TorchScript?

Onafhankelijk van Python
Efficiënt in productie
Voorbeelden:
- Uitrol op mobiele apparaten
- Snelle inferentie in productiesystemen

Een mobiele telefoon en een laptop met tandwielen die een productie-omgeving voorstellen

Modellen omzetten naar TorchScript

Twee conversiemethoden:

torch.jit.trace: volgt uitvoering met voorbeeldinputs
torch.jit.script: compileert door Python-code te analyseren

Wanneer gebruiken:

Gebruik trace voor eenvoudige modellen
Gebruik script voor modellen met controlflow (bijv. loops)

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-modellen opslaan en laden

Model opslaan:
- torch.jit.save: sla het gescripte model op in een bestand
Model laden:
- torch.jit.load: laad het model voor inference

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

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

Inferentie uitvoeren met TorchScript

Stappen:
- Laad het TorchScript-model
- Geef inputs door voor voorspellingen
- Outputs zijn gelijk aan PyTorch-voorspellingen

Voorbeeldinput:

Input-tensor: [1.0, 2.0, 3.0]

Voorbeeldoutput:

Output-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 in het kort

torch.jit.trace: werkt voor statische modellen
torch.jit.script: ondersteunt dynamische controlflow
torch.jit.save: slaat het gescripte model op
torch.jit.load: laadt opnieuw voor inferentie

Laten we oefenen!

Schaalbare AI-modellen met PyTorch Lightning