실습 5 — 지식증류: 큰 모델의 답으로 작은 모델 가르치기

import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
Xtr, Xte, ytr, yte = train_test_split(X / 16.0, y, test_size=0.3, random_state=42, stratify=y)
Xtr, Xte = torch.tensor(Xtr, dtype=torch.float32), torch.tensor(Xte, dtype=torch.float32)
ytr, yte = torch.tensor(ytr), torch.tensor(yte)


def mlp(h):
    return nn.Sequential(nn.Linear(64, h), nn.ReLU(), nn.Linear(h, h), nn.ReLU(), nn.Linear(h, 10))


def fit(model, X, target, epochs=60, seed=0, soft=False, T=4.0):
    """soft=False 면 정답 번호로, soft=True 면 선생님의 점수(로짓)를 온도 T 로 부드럽게 한 확률을 흉내 내게 학습"""
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    g = torch.Generator().manual_seed(seed)
    for _ in range(epochs):
        idx = torch.randperm(len(X), generator=g)
        for i in range(0, len(X), 32):
            b = idx[i:i + 32]
            out = model(X[b])
            if soft:   # 학생과 선생님의 부드러운 확률 사이의 거리(KL). T² 를 곱해 크기를 맞춘다
                loss = F.kl_div(F.log_softmax(out / T, 1), F.softmax(target[b] / T, 1), reduction="batchmean") * T * T
            else:
                loss = F.cross_entropy(out, target[b])
            opt.zero_grad()
            loss.backward()
            opt.step()
    return model


accuracy = lambda m: round((m(Xte).argmax(1) == yte).float().mean().item(), 3)
params = lambda m: sum(p.numel() for p in m.parameters())
torch.manual_seed(0)
teacher = fit(mlp(256), Xtr, ytr)
print("선생님: 파라미터", params(teacher), "· 정확도", accuracy(teacher), "· 학생 크기:", params(mlp(16)))

with torch.no_grad():
    logits = teacher(Xtr)                                  # 학습 그림 1,257장에 선생님이 매긴 점수
i = int((logits.softmax(1).max(1).values).argmin())        # 선생님이 가장 헷갈린 그림
for T in (1, 4):
    p = (logits[i] / T).softmax(0)
    top = p.topk(3)
    print(f"T={T}: 정답 {ytr[i].item()} · 선생님의 상위 셋", [(int(k), round(float(v), 3)) for v, k in zip(top.values, top.indices)])

n_labeled = 100                                            # 정답 표시가 있는 그림은 100장뿐이라고 가정
rows = []
torch.manual_seed(1)
s1 = fit(mlp(16), Xtr[:n_labeled], ytr[:n_labeled], epochs=200, seed=1)
rows.append(("정답 100개로만", accuracy(s1)))
torch.manual_seed(1)
s2 = fit(mlp(16), Xtr, logits.argmax(1), epochs=100, seed=1)              # 1,257장에 선생님이 고른 답(하드)
rows.append(("선생님의 답(하드)으로", accuracy(s2)))
torch.manual_seed(1)
s3 = fit(mlp(16), Xtr, logits, epochs=100, seed=1, soft=True)            # 1,257장에 선생님의 확률(소프트)
rows.append(("선생님의 확률(KD)로", accuracy(s3)))
torch.manual_seed(1)
s4 = fit(mlp(16), Xtr, ytr, epochs=100, seed=1)                           # 견줄 기준: 1,257장 모두 정답이 있다면
rows.append(("(참고) 정답 1,257개로", accuracy(s4)))
for name, acc in rows:
    print(f"{name}: {acc}")

def speed_ms(m):
    big_batch = Xte.repeat(50, 1)
    with torch.no_grad():
        t = time.perf_counter()
        for _ in range(5):
            m(big_batch)
    return round((time.perf_counter() - t) / 5 * 1000, 2)


for name, m in [("선생님", teacher), ("학생(KD)", s3)]:
    print(f"{name}: 파라미터 {params(m):,} · fp32 로 {params(m) * 4 / 1024:.1f}KB · 27,000개 추론 {speed_ms(m)}ms · 정확도 {accuracy(m)}")