Kendala umum saat melatih model deep learning menggunakan PyTorch adalah CUDA out of memory (OOM) yang tidak terjadi pada iterasi pertama, melainkan muncul secara bertahap setelah beberapa batch atau epoch. Pola ini mengindikasikan adanya memory leak struktural pada Python process atau VRAM GPU, bukan sekadar ukuran batch_size yang terlalu besar.

Root Cause: Akumulasi History Computation Graph

PyTorch menggunakan mekanisme define-by-run melalui modul Autograd. Setiap operasi tensor yang memiliki parameter requires_grad=True akan membangun DAG (Directed Acyclic Graph). Node grafik ini menyimpan referensi ke tensor input, gradien fungsi, dan konteks komputasi mundur.

Penyebab utama memory leak pada training loop adalah:

  • Menyimpan Tensor Loss Langsung: Variabel loss memegang referensi ke seluruh history komputasi forward pass batch tersebut. Jika diakumulasikan ke dalam variabel agregat Python (misal total_loss += loss atau history.append(loss)), garbage collector tidak dapat membebaskan memori graf tersebut. Graf terus tumbuh dari batch ke batch hingga VRAM habis.
  • Validasi Tanpa Context Manager Evaluasi: Menjalankan inference pada loop validasi tanpa torch.no_grad() tetap mengalokasikan memori untuk graf kalkulasi turunan yang sebenarnya tidak pernah digunakan oleh loss.backward().

Analisis Kode: Sebelum dan Sesudah Perbaikan

Kode di bawah menunjukkan implementasi umum yang menyebabkan memory leak bertahap:

# ANTI-PATTERN: Menyebabkan akumulasi graf komputasi
def train_bad(model, dataloader, optimizer, criterion):
    model.train()
    total_loss = 0.0
    loss_history = []

    for inputs, targets in dataloader:
        inputs, targets = inputs.cuda(), targets.cuda()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        # BUG 1: Menyimpan tensor dengan referensi ke graf backward pass
        total_loss += loss
        loss_history.append(loss)

def validate_bad(model, dataloader, criterion):
    # BUG 2: Tidak mematikan autograd dan tidak beralih ke eval mode
    val_loss = 0.0
    for inputs, targets in dataloader:
        inputs, targets = inputs.cuda(), targets.cuda()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        val_loss += loss

Solusi yang benar membutuhkan ekstraksi nilai skalar Python murni menggunakan .item() dan context manager torch.no_grad():

# BEST PRACTICE: Graf dibersihkan setiap iterasi
def train_good(model, dataloader, optimizer, criterion):
    model.train()
    total_loss = 0.0
    loss_history = []

    for inputs, targets in dataloader:
        inputs, targets = inputs.cuda(), targets.cuda()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        # FIX 1: Konversi tensor ke float Python primitif
        loss_scalar = loss.item()
        total_loss += loss_scalar
        loss_history.append(loss_scalar)

@torch.no_grad()
def validate_good(model, dataloader, criterion):
    # FIX 2: Set model ke eval mode dan nonaktifkan kalkulasi gradien
    model.eval()
    val_loss = 0.0
    for inputs, targets in dataloader:
        inputs, targets = inputs.cuda(), targets.cuda()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        val_loss += loss.item()

Verifikasi Memori dengan torch.cuda

Untuk memastikan apakah memori stabil sepanjang epoch, lakukan instrumentasi pengukuran memori pada akhir setiap batch menggunakan torch.cuda.memory_allocated():

import torch

def verify_gpu_memory(step: int):
    allocated = torch.cuda.memory_allocated() / (1024 ** 2)
    reserved = torch.cuda.memory_reserved() / (1024 ** 2)
    print(f"[Step {step}] Allocated: {allocated:.2f} MB | Reserved: {reserved:.2f} MB")

# Eksekusi loop training
for step, (x, y) in enumerate(dataloader):
    # Training step
    # ...
    if step % 50 == 0:
        verify_gpu_memory(step)

Jika nilai Allocated meningkat secara linier setiap 50 step, dipastikan terdapat pointer tensor atau autograd node yang tertahan di memori.

Checklist Preventif Training Loop PyTorch

  • Gunakan loss.item() atau float(loss) ketika mencatat loss ke log, tensorboard, atau agregasi metrik.
  • Hindari append tensor CUDA langsung ke list Python tanpa memanggil .detach().cpu() jika struktur array diperlukan untuk kalkulasi metrik lanjutan (seperti ROC-AUC atau F1-score).
  • Bungkus seluruh pipeline validasi dan inference dengan with torch.no_grad(): atau decorator @torch.no_grad().
  • Panggil model.eval() untuk menonaktifkan behavior training-layer seperti Dropout dan Batch Normalization.
  • Pastikan tidak menyimpan representasi output terprediksi berukuran penuh dalam memori jika hanya nilai argmax yang dibutuhkan.