Precizie Mixtă Automată (torch.autocast & GradScaler) în PyTorch & Deep Learning Practic

Accelerarea antrenării cu până la 2x-3x și reducerea consumului de memorie VRAM la jumătate folosind calcule rapide în float16/bfloat16 prin AMP.

Sintaxă & invocare Python
from torch.amp import autocast, GradScaler

scaler = GradScaler("cuda")

for batch_x, batch_y in loader:
    optimizator.zero_grad(set_to_none=True)
    # Calcule rapide în precizie mixtă (float16 / bfloat16):
    with autocast(device_type="cuda", dtype=torch.bfloat16):
        predictii = model(batch_x)
        loss = criteriu(predictii, batch_y)
    
    # Scalare automată împotriva dispariției gradienților:
    scaler.scale(loss).backward()
    scaler.step(optimizator)
    scaler.update()

Descriere & Utilizare

În mod tradițional, toate calculele din rețelele neuronale se desfășoară în format float32 (32 de biți). Însă plăcile video moderne (NVIDIA Tensor Cores de la arhitectura Volta până la Blackwell) au unități specializate de calcul capabile să execute operațiuni matriceale de două până la patru ori mai rapid în format de precizie redusă de 16 biți (float16 sau bfloat16).

AMP (Automatic Mixed Precision) combină cele mai bune aspecte din ambele lumi: torch.autocast alege automat cele mai potrivite operațiuni pentru 16 biți (precum convoluțiile și înmulțirile de matrice), păstrând operațiunile sensibile (precum Softmax și adunările de gradienți) în 32 de biți pentru a preveni erorile de rotunjire.

Dacă folosești float16, GradScaler scalează gradienții cu un factor mare pentru a preveni fenomenul de Underflow (când un gradient devine mai mic de 1e-7 și se transformă în zero curat).

Cum să o reții ușor

💡 Este ca desenarea unei schițe rapide în creion pe hârtie milimetrică (16 biți pentru 95% din schiță), urmată de trasarea conturului critic cu stilou de mare precizie (32 biți pentru nodurile cheie). Obții desenul de trei ori mai repede cu aceeași calitate finală.

Opțiuni și parametri tehnici

  • torch.autocast(device_type="cuda", dtype=torch.bfloat16): activează precizia mixtă în blocul de cod asociat.
  • torch.bfloat16: formatul recomandat pe GPU-urile moderne (NVIDIA Ampere RTX 30xx+, Hopper, Blackwell; Apple Silicon M-series); are același interval dinamic ca float32 și nu necesită GradScaler!
  • GradScaler("cuda"): necesar doar dacă rulezi pe plăci mai vechi în torch.float16.

Exemple practice și cazuri de utilizare

  • Antrenare modernă pe plăci RTX 4090 / A100 cu bfloat16: with autocast("cuda", dtype=torch.bfloat16): loss = model(x); loss.backward(); optimizer.step() (fără scaler!).
  • Dublarea mărimii de batch: activarea AMP îți permite de obicei să crești batch_size de la 16 la 32 pe același GPU.
  • Inferență accelerată: folosește autocast și în timpul inferenței pentru un răspuns de două ori mai rapid.

Sfaturi de configurare & capcane

Dacă GPU-ul tău suportă bfloat16 (verifică cu torch.cuda.is_bf16_supported()), folosește întotdeauna bfloat16 în loc de float16 — este mult mai stabil numeric și nu suferă niciodată de explozia sau dispariția gradienților.

Comenzi și ghiduri înrudite