torch.export: exportul modelelor în locul TorchScript în PyTorch & Deep Learning Practic

import torch
from torch.export import export, Dim

ep = export(model.eval(), args=(torch.randn(4, 16),),
            dynamic_shapes={"x": {0: Dim("batch")}})
torch.export.save(ep, "model.pt2")

ep2 = torch.export.load("model.pt2")
y = ep2.module()(torch.randn(9, 16))

torch.export.export() primește un nn.Module și intrări de exemplu și produce, ahead-of-time, un ExportedProgram: un graf care conține doar calculul cu tensori, la nivel de operatori ATen. Spre deosebire de torch.compile (compilator JIT care, la cod netrasabil, face „graph break” și revine la Python), torch.export cere graful complet și dă eroare la cod netrasabil. În schimb, graful rezultat nu mai depinde de Python și poate fi salvat, încărcat și rulat în alte medii.

TorchScript (torch.jit.script, torch.jit.trace) este depreciat. Notele de lansare 2.10 au adăugat avertismente de depreciere și spun că torch.jit nu este garantat pe Python 3.14; în 2.14 apelurile emit FutureWarning vizibil și trimit spre torch.compile sau torch.export. Documentația precizează diferența: torch.export nu capturează fluxul de control Python decât prin operatori expliciți (de exemplu torch.cond), dar acoperă mai multe construcții Python decât torch.jit.script.

Implicit, formele se specializează pe intrarea de exemplu: un program exportat cu lotul 3 refuză lotul 7. Dimensiunile dinamice se declară prin dynamic_shapes, cu Dim("nume", min=..., max=...), Dim.DYNAMIC sau Dim.AUTO. torch.export.save() / torch.export.load() folosesc o arhivă zip în formatul PT2 Archive. Pentru inferență fără Python pe server, AOTInductor compilează programul exportat cu torch._inductor.aoti_compile_and_package() într-un pachet .pt2, încărcat apoi din Python (aoti_load_package) sau din C++.

Concepte cheie și caracteristici

Concept / Parametru	Descriere practică
export(mod, args, kwargs=None, *, dynamic_shapes=None, strict=False)	Capturează modulul. args este un tuplu de intrări de exemplu.
dynamic_shapes	Dicționar pe numele argumentelor din forward: {"x": {0: Dim("batch")}}.
Dim("nume", min, max) / Dim.DYNAMIC / Dim.AUTO	Dimensiune simbolică; Dim.AUTO nu dă eroare dacă dimensiunea se dovedește statică.
ExportedProgram.module()	Întoarce un nn.Module care rulează graful exportat.
run_decompositions()	Descompune graful într-un set de operatori fără mutații și aliasing.
torch.export.save(ep, f) / torch.export.load(f)	Salvare și încărcare în formatul PT2 Archive (.pt2).
torch._inductor.aoti_compile_and_package(ep, package_path=...)	AOTInductor: compilează programul exportat pentru rulare fără Python.

Exemple practice și scenarii reale

Exemplul 1: Export cu lot dinamic, salvare și încărcare

import torch
import torch.nn as nn
from torch.export import export, Dim

class Clasificator(nn.Module):
    def __init__(self):
        super().__init__()
        self.strat = nn.Linear(16, 3)

    def forward(self, x):
        return torch.relu(self.strat(x))

ep = export(Clasificator().eval(), (torch.randn(4, 16),),
            dynamic_shapes={"x": {0: Dim("batch")}})
torch.export.save(ep, "clasificator.pt2")

ep2 = torch.export.load("clasificator.pt2")
print(ep2.module()(torch.randn(9, 16)).shape)   # torch.Size([9, 3])

Exemplul 2: Fără dynamic_shapes, forma rămâne fixă

ep_fix = export(Clasificator().eval(), (torch.randn(4, 16),))
ep_fix.module()(torch.randn(9, 16))   # eroare: forma a fost specializată la lotul 4

Exemplul 3: AOTInductor: pachet compilat pentru inferență

După documentația AOTInductor. Pe o mașină cu PyTorch compilat pentru CUDA pachetul rulează pe GPU, altfel pe CPU.

cale = torch._inductor.aoti_compile_and_package(ep, package_path="clasificator_aoti.pt2")
model_compilat = torch._inductor.aoti_load_package(cale)
print(model_compilat(torch.randn(5, 16)))

Greșeli frecvente și capcane

• Să pornești un proiect nou pe torch.jit.script sau torch.jit.trace. API-ul este depreciat; documentația trimite la torch.export (servire, alte limbaje) sau la torch.compile (viteză în Python).
• Să exporți fără dynamic_shapes un model care va primi loturi sau lungimi diferite. Programul refuză orice altă formă decât cea de exemplu.
• Să aștepți ca un if Python care depinde de valorile tensorilor să intre în graf. Exportul dă eroare de flux de control dependent de date; se rescrie cu torch.cond sau torch.where.

Noțiuni și module înrudite

• torch-compile
• torch-save-load-state-dict
• torch-nn-module