Menyimpan dan memuat

Hasil pelatihan berada di memori dan hilang begitu proses Python berakhir. Topik ini membahas cara memindahkan bobot model ke file dan mengembalikannya, mulai dari isi state_dict, penyimpanan dengan torch.save, pemuatan dengan load_state_dict, hingga pola checkpoint pelatihan lengkap yang menyertakan optimizer dan epoch.

Isi state_dict#

state_dict adalah kamus bertipe OrderedDict yang memetakan nama parameter ke tensor bobotnya. Nama kunci mengikuti struktur modul, yaitu indeks submodul untuk model nn.Sequential dan nama atribut untuk modul yang ditulis dengan kelas. Karena kuncinya stabil dan isinya tensor biasa, state_dict adalah bentuk yang tepat untuk disimpan dan dipindahkan antara program. Skrip berikut memeriksa isi state_dict pada model kecil:

"""Isi state_dict pada model kecil: nama kunci dan bentuk tensor-nya."""

import torch.nn as nn

model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 3))
state = model.state_dict()

print("tipe state_dict :", type(state).__name__)
for kunci, nilai in state.items():
    print(f"{kunci:10s} {tuple(nilai.shape)}")
print("jumlah kunci :", len(state))

tipe state_dict : OrderedDict
0.weight   (16, 2)
0.bias     (16,)
2.weight   (3, 16)
2.bias     (3,)
jumlah kunci : 4

Kunci 0.weight dan 2.bias berasal dari posisi modul di dalam nn.Sequential, sedangkan ReLU tidak muncul sama sekali karena tidak memiliki parameter. Angka indeks itu penting saat memuat, karena load_state_dict mencocokkan tensor berdasarkan kunci dan bentuknya. Apabila urutan atau arsitektur layer berubah, kunci tidak cocok dan pemuatan ditolak, sifat yang sekilas mengganggu tetapi melindungi dari pemuatan yang keliru.

Menyimpan dan memuat kembali#

Penyimpanan menggunakan torch.save(model.state_dict(), "model.pt"), dan pemuatan dilakukan pada model baru yang arsitekturnya sama. Method load_state_dict menyalin isi tensor ke parameter model sasaran, sehingga model baru berperilaku persis sama dengan model asal tanpa melatih ulang. Skrip berikut melatih model mini, menyimpan state_dict, memuatnya ke model baru, lalu membandingkan keluaran ketiga model pada masukan yang sama:

"""Melatih model mini, menyimpan state_dict, lalu memuatnya kembali."""

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)


def buat_model():
    """Model mini dengan arsitektur tetap, dipakai dua kali."""
    return nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 3))


# data sintetis: tiga gugus dengan dua fitur
n_per_kelas = 100
pusat = torch.tensor([[0.0, 0.0], [3.0, 3.0], [-3.0, 3.0]])
fitur = torch.cat([pusat[k] + 0.6 * torch.randn(n_per_kelas, 2) for k in range(3)])
label = torch.cat([torch.full((n_per_kelas,), k) for k in range(3)])
dataset = TensorDataset(fitur, label)
loader = DataLoader(dataset, batch_size=32, shuffle=True,
                    generator=torch.Generator().manual_seed(0))

model = buat_model()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
criterion = nn.CrossEntropyLoss()

# pelatihan singkat, 5 epoch
model.train()
for epoch in range(5):
    for xb, yb in loader:
        optimizer.zero_grad()
        loss = criterion(model(xb), yb)
        loss.backward()
        optimizer.step()

# simpan hanya state_dict, bukan seluruh objek model
torch.save(model.state_dict(), "model.pt")

# model baru dengan arsitektur yang sama, lalu muat bobotnya
model_muat = buat_model()
model_muat.load_state_dict(torch.load("model.pt", weights_only=True))

# bandingkan keluaran ketiga model pada masukan yang sama
model_seg = buat_model()
with torch.no_grad():
    y_asal = model(fitur)
    y_muat = model_muat(fitur)
    y_seg = model_seg(fitur)

print("allclose(model, model_muat):", torch.allclose(y_asal, y_muat))
print("allclose(model, model_seg) :", torch.allclose(y_asal, y_seg))

allclose(model, model_muat): True
allclose(model, model_seg) : False

Nilai True pada baris pertama berarti keluaran model asal dan model muat identik untuk setiap elemen, diperiksa dengan torch.allclose. Nilai False pada baris kedua berasal dari model yang baru dibuat tanpa pemuatan, yang bobotnya berbeda karena inisialisasi acak. Perbandingan keluaran seperti ini adalah cara paling langsung untuk memastikan pemuatan berhasil, jauh sebelum model digunakan pada data sungguhan. Alur lengkapnya digambarkan sebagai berikut:

Checkpoint pelatihan lengkap#

state_dict model saja cukup untuk inferensi, tetapi tidak cukup untuk melanjutkan pelatihan yang terputus. Optimizer menyimpan keadaan internal seperti momentum pada Adam, dan pelatihan yang dilanjutkan perlu mengetahui epoch terakhir yang tercapai. Checkpoint lengkap menyimpan semuanya dalam satu kamus:

# tidak-dijalankan
checkpoint = {
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "val_loss": val_loss,
}
torch.save(checkpoint, "checkpoint.pt")

ckpt = torch.load("checkpoint.pt", weights_only=True)
model.load_state_dict(ckpt["model_state_dict"])
optimizer.load_state_dict(ckpt["optimizer_state_dict"])

Pola di atas tidak dijalankan sebagai file utuh karena epoch, model, optimizer, dan val_loss berasal dari perulangan pelatihan yang membungkusnya. Bentuknya tetap sah untuk disimpan, karena kamus yang berisi campuran tensor dan bilangan biasa tetap dapat dimuat dengan weights_only=True. Untuk melanjutkan pelatihan, model dan optimizer dibuat ulang terlebih dahulu, lalu kedua state_dict dimuat dari kamus yang sama sebelum perulangan epoch dilanjutkan dari nilai epoch yang tersimpan.

Keamanan torch.load dan weights_only#

torch.save menulis file dengan format pickle, yaitu format serialisasi Python yang mengeksekusi instruksi tersimpan saat file dimuat. File pickle dari sumber yang tidak dikenal dapat menjalankan kode apa pun di mesin pembacanya, sehingga pemuatan adalah titik rawan keamanan pada alur kerja model.

Argumen weights_only=True membatasi pemuatan hanya pada tensor dan tipe data primitif, sehingga instruksi berbahaya tidak tereksekusi. PyTorch 2.6 dan seterusnya, termasuk 2.15 yang digunakan pada jalur ini, menjadikan nilai tersebut default untuk torch.load. Penulisan eksplisit tetap dianjurkan agar maksudnya terbaca langsung pada kode dan kebiasaannya terbawa ke versi yang lebih lama. Seluruh contoh di topik ini menggunakan bentuk eksplisit torch.load("model.pt", weights_only=True).

Kesalahan umum#

Kesalahan pertama adalah memuat state_dict ke arsitektur yang berbeda. load_state_dict menolak dengan RuntimeError yang menyebut kunci yang hilang, kunci yang tak terduga, atau ketidakcocokan bentuk. Perbaikannya adalah memastikan definisi model sama persis dengan model asal, dan argumen strict=False hanya digunakan bila pemuatan sebagian memang disengaja.

Kesalahan kedua adalah menyimpan seluruh objek model dengan torch.save(model). File hasilnya terikat pada definisi kelas dan struktur direktori saat disimpan, sehingga rawan gagal dimuat ketika dipindahkan ke proyek lain. Perbaikannya adalah menyimpan model.state_dict() sebagaimana seluruh contoh di topik ini.

Kesalahan ketiga adalah lupa memanggil model.eval() setelah memuat bobot untuk inferensi. Dropout yang masih aktif membuat keluaran berubah pada setiap eksekusi meskipun bobotnya sama. Perbaikannya adalah menjadikan load_state_dict dan model.eval() sebagai satu rangkaian yang tidak terpisah pada tahap inferensi.

Ringkasan#

• state_dict memetakan nama parameter ke tensor, dan bentuk itulah yang disimpan, bukan seluruh objek model.
• torch.save menulis state_dict ke file, dan load_state_dict mengembalikannya ke model baru dengan arsitektur yang sama.
• Checkpoint pelatihan lengkap menyertakan optimizer, epoch, dan loss validasi agar pelatihan dapat dilanjutkan.
• torch.load dengan weights_only=True adalah kebiasaan aman untuk file dari sumber mana pun.

Referensi#

• Saving and Loading Models — PyTorch — Tutorial resmi tentang state_dict, checkpoint, dan pemuatan yang aman.
• PyTorch documentation — PyTorch — Rujukan API resmi untuk torch.save, torch.load, dan load_state_dict.