Salvarea și încărcarea modelelor: state_dict, checkpoint-uri, weights_only în PyTorch: sintaxă Python, tensori și exemple practice

# Salvare: doar parametrii și bufferele torch.save(model.state_dict(), "model.pt") # Încărcare într-un model cu aceeași arhitectură model = ReteaClasificare(784, 10) model.load_state_dict(torch.load("model.pt", weights_only=True, map_location="cpu")) model.eval()

model.state_dict() este un dicționar care leagă numele fiecărui parametru și buffer de tensorul lui. Documentația despre serializare recomandă salvarea lui, nu a modulului întreg: torch.save(model) folosește pickle pe obiect și depinde de clasa și de calea ei din cod. La încărcare creezi întâi modelul, apoi apelezi load_state_dict().

torch.load folosește un unpickler. Începând cu PyTorch 2.6, weights_only=True este valoarea implicită (schimbare anunțată prin avertisment din 2.4): se pot încărca doar tensori, tipuri primitive, dicționare și tipurile permise explicit cu torch.serialization.add_safe_globals(). Un fișier salvat cu torch.save(model) dă acum UnpicklingError („Weights only load failed”); documentația spune că weights_only=False e acceptabil doar pentru fișiere în care ai încredere și avertizează să nu încarci niciodată date din surse nesigure. weights_only=True reduce suprafața de atac, dar nu protejează de atacuri de tip denial of service.

Un checkpoint de reluare conține, pe lângă model, starea optimizatorului, a planificatorului și epoca. map_location mută tensorii la încărcare (de exemplu pe CPU, dacă au fost salvați pe un GPU care lipsește), iar mmap=True mapează fișierul în loc să-l citească tot în memorie. Pentru modele distribuite pe mai multe GPU-uri există Distributed Checkpoint (DCP), iar pentru rulare fără codul Python al modelului, torch.export.

Concepte cheie și caracteristici

Concept / Parametru	Descriere practică
`torch.save(obj, f)`	Serializează un obiect (de obicei un dicționar cu state_dict-uri).
`torch.load(f, map_location=None, *, weights_only=True, mmap=None)`	Încarcă; `weights_only=True` implicit din 2.6.
`map_location`	`"cpu"`, un `torch.device`, un dicționar de remapare sau o funcție.
`model.load_state_dict(sd, strict=True)`	Copiază valorile; `strict=False` tolerează chei lipsă sau în plus și le raportează.
`torch.serialization.add_safe_globals([...])`	Permite tipuri suplimentare la `weights_only=True`.
`mmap=True`	Încarcă leneș tensorii din fișier (mai puțină memorie RAM).

Exemple practice și scenarii reale

Exemplul 1: Checkpoint complet și reluare

torch.save({
    "epoca": epoca,
    "model": model.state_dict(),
    "optimizator": optimizator.state_dict(),
    "scheduler": scheduler.state_dict(),
    "pierdere": cel_mai_bun,
}, "checkpoint.pt")

ck = torch.load("checkpoint.pt", weights_only=True, map_location=device)
model.load_state_dict(ck["model"])
optimizator.load_state_dict(ck["optimizator"])
scheduler.load_state_dict(ck["scheduler"])
epoca_start = ck["epoca"] + 1

Exemplul 2: De ce eșuează un model salvat întreg

torch.save(model, "tot.pt")
torch.load("tot.pt")   # UnpicklingError: Weights only load failed ...
# Doar pentru fișiere proprii, de încredere:
model = torch.load("tot.pt", weights_only=False)

Exemplul 3: Format safetensors

Biblioteca safetensors (Hugging Face, instalată separat) salvează tensorii fără pickle.

from safetensors.torch import save_file, load_file

save_file(model.state_dict(), "model.safetensors")
model.load_state_dict(load_file("model.safetensors"))

Greșeli frecvente și capcane

Să încarci cu weights_only=False un fișier descărcat de pe internet. Unpickler-ul poate executa cod la încărcare.
Să salvezi doar modelul când vrei să reiei antrenarea. Fără starea optimizatorului (momentele lui Adam) și a planificatorului, reluarea nu continuă de unde ai rămas.
Să încarci un state_dict salvat dintr-un model împachetat în DDP direct în modelul simplu. Cheile au prefixul module.; salvează model.module.state_dict() (vezi DDP).

`Salvarea și încărcarea modelelor: state_dict, checkpoint-uri, weights_only` în PyTorch & Deep Learning Practic