실습 4 — 프루닝: 가지를 쳐도 맞히나

import copy, time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.prune as prune
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False

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, y, epochs=60, seed=0):
    """작은 묶음(32개)으로 Adam 학습. 씨앗을 고정해 다시 돌려도 같다"""
    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]
            loss = F.cross_entropy(model(X[b]), y[b])
            opt.zero_grad()
            loss.backward()
            opt.step()
    return model


accuracy = lambda m: round((m(Xte).argmax(1) == yte).float().mean().item(), 3)
torch.manual_seed(0)
dense = fit(mlp(256), Xtr, ytr)
print("파라미터:", sum(p.numel() for p in dense.parameters()), "· 시험 정확도:", accuracy(dense))

def pruned_copy(model, amount):
    """모든 층의 가중치를 한데 모아, 절댓값이 작은 쪽부터 amount 만큼 0 으로 만든 복사본"""
    m = copy.deepcopy(model)
    params = [(m[0], "weight"), (m[2], "weight"), (m[4], "weight")]
    if amount > 0:
        prune.global_unstructured(params, pruning_method=prune.L1Unstructured, amount=amount)
    return m


amounts = [0, 0.3, 0.5, 0.7, 0.8, 0.9, 0.95]
curve = [accuracy(pruned_copy(dense, a)) for a in amounts]
print(dict(zip(amounts, curve)))
plt.plot([a * 100 for a in amounts], curve, "o-")
plt.xlabel("지운 가중치의 몫(%)")
plt.ylabel("시험 정확도")
plt.title("비정형 프루닝: 지울수록 정확도는?")
plt.grid(alpha=0.3)
plt.show()

m90 = pruned_copy(dense, 0.9)
before = accuracy(m90)
fit(m90, Xtr, ytr, epochs=20, seed=1)              # 지운 자리는 마스크로 0 에 묶인 채 남은 가중치만 다시 학습
zeros = sum((layer.weight == 0).sum().item() for layer in (m90[0], m90[2], m90[4]))
total = sum(layer.weight.numel() for layer in (m90[0], m90[2], m90[4]))
print("다시 학습 전", before, "→ 뒤", accuracy(m90), f"· 0 인 가중치 {zeros / total:.1%}")

def structured(model, keep):
    """숨은층마다 가중치 크기(L1)가 큰 뉴런 keep 개만 남긴 더 작은 신경망을 만든다"""
    W0, W2, W4 = model[0].weight.data, model[2].weight.data, model[4].weight.data
    k1 = W0.abs().sum(1).topk(keep).indices                      # 첫 숨은층에서 남길 뉴런
    k2 = W2[:, k1].abs().sum(1).topk(keep).indices               # 둘째 숨은층에서 남길 뉴런
    small = mlp(keep)
    small[0].weight.data, small[0].bias.data = W0[k1].clone(), model[0].bias.data[k1].clone()
    small[2].weight.data, small[2].bias.data = W2[k2][:, k1].clone(), model[2].bias.data[k2].clone()
    small[4].weight.data, small[4].bias.data = W4[:, k2].clone(), model[4].bias.data.clone()
    return small


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


rows = []
for keep in (256, 128, 64, 32):
    s = structured(dense, keep)
    a0 = accuracy(s)
    fit(s, Xtr, ytr, epochs=20, seed=1)
    rows.append({"남긴 뉴런": keep, "파라미터": sum(p.numel() for p in s.parameters()),
                 "지운 직후": a0, "다시 학습 뒤": accuracy(s), "27,000개 추론 ms": speed_ms(s)})
import pandas as pd
print(pd.DataFrame(rows).to_string(index=False))