Lab 6: 적대적 훈련으로 모델 방어하기

pip install adversarial-robustness-toolbox torch torchvision numpy

python -c "import art; print(art.__version__)"
python -c "import torch; print(torch.__version__)"

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

transform = transforms.Compose([transforms.ToTensor()])
train_set = torchvision.datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_set = torchvision.datasets.MNIST(root="./data", train=False, download=True, transform=transform)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
test_loader = DataLoader(test_set, batch_size=256, shuffle=False)

def train(model, loader, epochs=2):
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()
    model.train()
    for epoch in range(epochs):
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            loss = criterion(model(x), y)
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch+1} 완료")
    return optimizer, criterion

baseline_model = SimpleCNN().to(device)
optimizer, criterion = train(baseline_model, train_loader, epochs=2)

x_test = np.concatenate([x.numpy() for x, _ in test_loader])
y_test = np.concatenate([y.numpy() for _, y in test_loader])

from art.estimators.classification import PyTorchClassifier
from art.attacks.evasion import FastGradientMethod, ProjectedGradientDescent

baseline_classifier = PyTorchClassifier(
    model=baseline_model,
    loss=criterion,
    optimizer=optimizer,
    input_shape=(1, 28, 28),
    nb_classes=10,
    clip_values=(0.0, 1.0),
    device_type="gpu" if torch.cuda.is_available() else "cpu",
)

def evaluate(classifier, x, y):
    preds = np.argmax(classifier.predict(x), axis=1)
    return np.mean(preds == y)

# Clean accuracy
acc_clean_before = evaluate(baseline_classifier, x_test, y_test)

# FGSM robust accuracy
fgsm = FastGradientMethod(estimator=baseline_classifier, eps=0.2)
x_test_fgsm = fgsm.generate(x=x_test)
acc_fgsm_before = evaluate(baseline_classifier, x_test_fgsm, y_test)

# PGD robust accuracy
pgd_eval = ProjectedGradientDescent(estimator=baseline_classifier, eps=0.2, eps_step=0.02, max_iter=40)
x_test_pgd = pgd_eval.generate(x=x_test)
acc_pgd_before = evaluate(baseline_classifier, x_test_pgd, y_test)

print(f"[방어 전] Clean accuracy: {acc_clean_before * 100:.2f}%")
print(f"[방어 전] FGSM(eps=0.2) robust accuracy: {acc_fgsm_before * 100:.2f}%")
print(f"[방어 전] PGD(eps=0.2, max_iter=40) robust accuracy: {acc_pgd_before * 100:.2f}%")

adv_model = SimpleCNN().to(device)
adv_optimizer = optim.Adam(adv_model.parameters(), lr=1e-3)
adv_criterion = nn.CrossEntropyLoss()

adv_classifier = PyTorchClassifier(
    model=adv_model,
    loss=adv_criterion,
    optimizer=adv_optimizer,
    input_shape=(1, 28, 28),
    nb_classes=10,
    clip_values=(0.0, 1.0),
    device_type="gpu" if torch.cuda.is_available() else "cpu",
)

# 훈련용 PGD: 빠른 훈련을 위해 max_iter를 평가용보다 작게 설정 (대표적인 비용 절감 패턴)
pgd_train = ProjectedGradientDescent(estimator=adv_classifier, eps=0.2, eps_step=0.04, max_iter=7)

def adversarial_train(classifier, attack, loader, epochs=2):
    for epoch in range(epochs):
        for x, y in loader:
            x_np, y_np = x.numpy(), y.numpy()
            # 미니배치 절반은 깨끗한 예제, 절반은 PGD 적대적 예제로 구성
            half = len(x_np) // 2
            x_adv = attack.generate(x=x_np[:half])
            x_mixed = np.concatenate([x_adv, x_np[half:]], axis=0)
            y_mixed = y_np
            classifier.fit(x_mixed, y_mixed, batch_size=len(x_mixed), nb_epochs=1, verbose=False)
        print(f"Adversarial training epoch {epoch+1} 완료")

adversarial_train(adv_classifier, pgd_train, train_loader, epochs=2)

# Clean accuracy (방어 후)
acc_clean_after = evaluate(adv_classifier, x_test, y_test)

# FGSM robust accuracy (방어 후)
fgsm_after = FastGradientMethod(estimator=adv_classifier, eps=0.2)
x_test_fgsm_after = fgsm_after.generate(x=x_test)
acc_fgsm_after = evaluate(adv_classifier, x_test_fgsm_after, y_test)

# PGD robust accuracy (방어 후)
pgd_after = ProjectedGradientDescent(estimator=adv_classifier, eps=0.2, eps_step=0.02, max_iter=40)
x_test_pgd_after = pgd_after.generate(x=x_test)
acc_pgd_after = evaluate(adv_classifier, x_test_pgd_after, y_test)

print(f"[방어 후] Clean accuracy: {acc_clean_after * 100:.2f}%")
print(f"[방어 후] FGSM(eps=0.2) robust accuracy: {acc_fgsm_after * 100:.2f}%")
print(f"[방어 후] PGD(eps=0.2, max_iter=40) robust accuracy: {acc_pgd_after * 100:.2f}%")

print("\n=== 방어 전/후 비교 ===")
print(f"Clean accuracy:        {acc_clean_before*100:.2f}% -> {acc_clean_after*100:.2f}%")
print(f"FGSM robust accuracy:  {acc_fgsm_before*100:.2f}% -> {acc_fgsm_after*100:.2f}%")
print(f"PGD robust accuracy:   {acc_pgd_before*100:.2f}% -> {acc_pgd_after*100:.2f}%")