Pembekuan dan penggantian head

ResNet18 yang dimuat dengan bobot pretrained selalu menghasilkan 1000 logits ImageNet, sedangkan tugas baru umumnya hanya memiliki beberapa kelas. Topik ini membahas dua langkah untuk menyiasatinya, yaitu mengganti head dengan layer klasifikasi baru dan membekukan backbone agar bobot pretrained tidak bergeser selama pelatihan. Kombinasi kedua langkah itu dikenal sebagai ekstraksi fitur (feature extraction), yaitu memanfaatkan model pretrained sebagai penghasil fitur yang tidak berubah sambil melatih head baru dari nol.

Anatomi backbone dan head#

Dalam konteks transfer learning, jaringan klasifikasi dibagi menjadi dua bagian. Backbone, yaitu rangkaian layer penghasil fitur, memetakan citra masukan menjadi sebuah vektor representasi. Head, yaitu layer klasifikasi di ujung jaringan, memetakan vektor fitur itu menjadi logits per kelas. Pembagian ini bukan konsep abstrak pada ResNet18, melainkan terbaca langsung dari daftar layer penyusunnya:

from torchvision.models import resnet18, ResNet18_Weights

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
for nama, layer in model.named_children():
    print(f"{nama:10} {layer.__class__.__name__}")

conv1      Conv2d
bn1        BatchNorm2d
relu       ReLU
maxpool    MaxPool2d
layer1     Sequential
layer2     Sequential
layer3     Sequential
layer4     Sequential
avgpool    AdaptiveAvgPool2d
fc         Linear

Keluaran menunjukkan sembilan layer backbone berurutan, dari konvolusi pembuka conv1 hingga layer4, ditutup pengambilan rata-rata global avgpool yang meratakan peta fitur menjadi vektor 512 nilai. Layer fc di baris terakhir adalah head-nya, yaitu Linear(512, 1000) yang menunjuk kelas ImageNet. Kedua bagian ini nantinya diperlakukan berbeda: head diganti karena menunjuk kelas ImageNet yang tidak sesuai tugas, sedangkan backbone dibekukan agar fitur hasil pelatihan ImageNet terjaga.

Mengganti head dengan layer baru#

Mengganti head pada ResNet18 semudah menimpa atribut fc dengan layer baru, sebab fc hanyalah nn.Module biasa yang dapat ditukar dengan modul lain berukuran masukan yang sama:

import torch.nn as nn
from torchvision.models import resnet18, ResNet18_Weights

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
print("head lama :", model.fc)
model.fc = nn.Linear(512, 4)
print("head baru :", model.fc)

head lama : Linear(in_features=512, out_features=1000, bias=True)
head baru : Linear(in_features=512, out_features=4, bias=True)

Argumen pertama nn.Linear harus 512 karena itulah lebar vektor fitur keluaran avgpool, sedangkan argumen kedua diisi jumlah kelas tugas baru, dalam contoh ini empat. Layer baru dibuat dengan inisialisasi acak dan berstatus requires_grad = True sebagaimana layer nn.Linear pada umumnya, sehingga siap dilatih. Perlu diketahui bahwa penimpaan membuang bobot head lama beserta 513000 parameternya, dan pembuangan itu permanen kecuali state_dict semula disimpan sebelumnya.

Membekukan backbone#

Membekukan sebuah parameter berarti menetapkan requires_grad = False, sehingga autograd tidak lagi menghitung dan mengakumulasikan gradien untuknya. Parameter beku tidak bergerak saat optimizer.step() dipanggil, sebab optimasi hanya bekerja pada parameter yang membawa gradien. Penetapan dapat dilakukan per parameter ataupun serentak melalui method requires_grad_() milik modul:

import torch.nn as nn

layer = nn.Linear(4, 2)
layer.weight.requires_grad_(False)
print("weight.requires_grad:", layer.weight.requires_grad)
print("bias.requires_grad  :", layer.bias.requires_grad)

weight.requires_grad: False
bias.requires_grad  : True

Keluaran menegaskan bahwa requires_grad_(False) hanya menyasar target yang ditunjuk, sehingga bias tetap dapat dilatih. Untuk membekukan backbone ResNet18, seluruh parameter dibekukan lebih dahulu, kemudian head baru dibuat. Urutan ini penting: head yang dibuat setelah pembekuan tidak tersentuh pembekuan itu dan tetap dapat dilatih. Skrip berikut merangkum urutan tersebut beserta hitungan parameternya:

import torch.nn as nn
from torchvision.models import resnet18, ResNet18_Weights

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

# bekukan backbone: seluruh bobot pretrained tidak menerima gradien lagi
for p in model.parameters():
    p.requires_grad = False

# head baru dibuat setelah pembekuan, sehingga tetap dilatih
model.fc = nn.Linear(512, 4)

total = sum(p.numel() for p in model.parameters())
dilatih = sum(p.numel() for p in model.parameters() if p.requires_grad)
print("parameter total  :", total)
print("parameter dilatih:", dilatih)
print("parameter beku   :", total - dilatih)
print("head baru        :", model.fc)

parameter total  : 11178564
parameter dilatih: 2052
parameter beku   : 11176512
head baru        : Linear(in_features=512, out_features=4, bias=True)

Dari 11178564 parameter, hanya 2052 milik head baru yang akan bergerak, sedangkan 11176512 parameter backbone terkunci pada bobot ImageNet-nya. Kondisi ini berpengaruh langsung pada optimizer: daftar parameter yang diserahkan kepadanya sebaiknya hanya berisi parameter yang dilatih, yang diperoleh dengan menyaring p.requires_grad. Pola lengkapnya digunakan pada skrip pelatihan di bagian akhir topik ini. Struktur model hasil kedua langkah itu digambarkan pada diagram berikut:

Ekstraksi fitur versus fine-tuning#

Membekukan seluruh backbone atau tidak membekukan apa pun sama-sama sah, dan pilihan itu membentuk dua strategi dengan tujuan berbeda. Fine-tuning, yaitu melatih ulang seluruh atau sebagian jaringan dengan learning rate kecil dari titik awal bobot pretrained, dibahas penuh pada topik Praktik fine-tuning. Tabel berikut menyandingkan keduanya.

Tabel itu menunjukkan bahwa ekstraksi fitur adalah titik awal yang murah dan aman, bukan sekadar langkah persiapan. Dengan data yang sedikit, kemungkinan overfit jauh lebih kecil ketika yang dilatih hanya 2052 parameter dibandingkan 11178564 parameter. Pilihan keduanya juga tidak biner, karena sebagian backbone dapat dibekukan sementara blok-blok akhir tetap dilatih, dan pengaturan semacam itu memerlukan learning rate yang dibedakan per kelompok parameter.

Melatih head pada data sintetis#

Efektivitas strategi ini dapat dibuktikan dengan data sintetis kecil. Data berupa 64 citra 64x64 yang kelasnya ditandai oleh kuadran tempat blok terang berada, ditambah 32 citra serupa untuk validasi, sehingga tidak ada unduhan sama sekali. Head empat kelas dilatih empat epoch dengan learning rate 0.02, dan layer batch normalization backbone dipertahankan pada mode eval selama pelatihan karena alasannya dijelaskan pada bagian Kesalahan umum:

import torch
import torch.nn as nn
from torchvision.models import resnet18, ResNet18_Weights

torch.manual_seed(0)


def data_sintetis(jumlah, ukuran=64):
    # kelas ditandai kuadran tempat blok terang berada, latar diberi derau
    x = torch.rand(jumlah, 3, ukuran, ukuran) * 0.15 + 0.35
    y = torch.zeros(jumlah, dtype=torch.long)
    setengah = ukuran // 2
    urut = torch.randperm(jumlah)
    for i, j in enumerate(urut):
        kelas = i % 4
        y[j] = kelas
        awal_h = 6 if kelas < 2 else setengah + 2
        awal_w = 6 if kelas % 2 == 0 else setengah + 2
        panjang = setengah // 2 - 2
        x[j, :, awal_h:awal_h + panjang, awal_w:awal_w + panjang] += 0.55
    return x, y


x_latih, y_latih = data_sintetis(64)
x_val, y_val = data_sintetis(32)
loss_fn = nn.CrossEntropyLoss()

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
model.requires_grad_(False)      # bekukan seluruh bobot pretrained
model.fc = nn.Linear(512, 4)     # head baru, otomatis requires_grad=True

param_dilatih = [p for p in model.parameters() if p.requires_grad]
print("tensor dilatih    :", len(param_dilatih), "dari", len(list(model.parameters())))
print("parameter dilatih :", sum(p.numel() for p in param_dilatih))

optimizer = torch.optim.SGD(param_dilatih, lr=0.02)
for epoch in range(4):
    model.train()
    for layer in model.modules():
        # statistik batch normalization backbone dipertahankan pada mode eval
        if isinstance(layer, nn.BatchNorm2d):
            layer.eval()
    total = 0.0
    for i in range(0, 64, 16):
        xb, yb = x_latih[i:i + 16], y_latih[i:i + 16]
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        total += loss.item() * len(xb)
    print(f"epoch {epoch + 1}: loss = {total / 64:.4f}")

model.eval()
with torch.no_grad():
    benar = (model(x_val).argmax(dim=1) == y_val).sum().item()
print(f"akurasi validasi  : {benar}/32")

tensor dilatih    : 2 dari 62
parameter dilatih : 2052
epoch 1: loss = 2.6960
epoch 2: loss = 2.2183
epoch 3: loss = 0.9061
epoch 4: loss = 0.4938
akurasi validasi  : 32/32

Keluaran mengonfirmasi klaim utama topik ini. Optimizer hanya menerima dua tensor milik head, yaitu bobot dan bias, dan hanya 2052 parameter yang bergerak, sedangkan loss turun dari 2.6960 menjadi 0.4938 dalam empat epoch. Akurasi validasi mencapai 32 dari 32 meskipun bobot backbone tidak tersentuh satu pun, karena fitur ImageNet yang dibekukan sudah memadai untuk memisahkan pola sederhana tersebut.

Kesalahan umum#

Membekukan seluruh parameter, termasuk head yang baru dibuat. Kesalahan ini terjadi apabila head dibuat lebih dahulu dan pembekuan dijalankan setelahnya, sehingga tidak tersisa satu pun parameter yang dapat dilatih. Gejalanya khas, yaitu loss tetap pada nilai awal tanpa pernah turun karena optimizer.step() tidak memiliki gradien untuk diterapkan. Perbaikannya adalah membekukan backbone sebelum membuat head, atau mengembalikan status head dengan model.fc.requires_grad_(True).

Mengira requires_grad = False membekukan segala sesuatu. Pembekuan itu hanya menghentikan gradien, sedangkan buffer statistik running_mean dan running_var milik batch normalization tetap diperbarui selama model berada di mode train. Data yang distribusinya jauh dari ImageNet pun menggeser statistik tersebut sehingga keluaran mode eval ikut rusak. Program berikut membuktikannya pada satu layer batch normalization:

import torch
import torch.nn as nn
from torchvision.models import resnet18, ResNet18_Weights

torch.manual_seed(0)
model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
for p in model.parameters():
    p.requires_grad = False

bn = model.layer1[0].bn1
print("running_var awal  :", bn.running_var[0].item())
model.train()                     # mode train, meski semua parameter beku
with torch.no_grad():
    model(torch.rand(4, 3, 64, 64))
print("running_var train :", bn.running_var[0].item())
model.eval()
with torch.no_grad():
    model(torch.rand(4, 3, 64, 64))
print("running_var eval  :", bn.running_var[0].item())

running_var awal  : 0.4350549876689911
running_var train : 0.4212796092033386
running_var eval  : 0.4212796092033386

Nilai running_var bergeser dari 0.4350 menjadi 0.4213 setelah satu forward pada mode train, meskipun seluruh parameter model beku, dan tidak berubah lagi saat mode eval. Inilah alasan skrip latih-head.py menempatkan semua layer batch normalization pada mode eval selama pelatihan head. Kebiasaan yang sama juga berlaku ketika data asli sedikit dan distribusinya berjauhan dari ImageNet.

Ringkasan#

• Backbone ResNet18 mencakup conv1 hingga avgpool yang menghasilkan vektor 512 fitur, sedangkan head-nya adalah fc dengan 1000 kelas ImageNet.
• Head diganti dengan menimpa model.fc menggunakan nn.Linear(512, jumlah_kelas), dan layer baru tersebut siap dilatih sejak awal.
• Membekukan parameter dengan requires_grad = False menghentikan penghitungan gradien, sehingga optimizer yang disaring dengan p.requires_grad hanya memperbarui head.
• requires_grad = False tidak menghentikan pembaruan statistik batch normalization pada mode train, sehingga layer tersebut perlu dipertahankan pada mode eval saat melatih head.

Referensi#

• Transfer Learning for Computer Vision Tutorial — PyTorch — Tutorial resmi yang mengganti head ResNet18 pretrained dan melatihnya untuk tugas dua kelas.
• Finetuning Torchvision Models — PyTorch — Pola feature extracting resmi yang membekukan parameter dan mengganti head.
• torch.nn.Module — PyTorch — Dokumentasi requires_grad_(), named_children(), dan modules() yang digunakan sepanjang topik ini.