8. 직접 해보기: 간단한 코드 예제 💻

이론만 배우는 것보다 직접 코드를 작성해보는 것이 이해에 도움이 되지! 여기서는 PyTorch를 사용해 간단한 CNN과 ViT 모델을 구현하는 방법을 알아볼게. 이 코드들은 기본적인 구조를 이해하는 데 도움이 될 거야! 🧑‍💻

간단한 CNN 모델 구현하기 🔧

먼저 PyTorch로 간단한 CNN 모델을 구현해볼게. 이 모델은 MNIST 손글씨 숫자 데이터셋을 분류하는 데 사용할 수 있어!

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 첫 번째 합성곱 층: 입력 채널 1, 출력 채널 32, 커널 크기 3x3
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        # 두 번째 합성곱 층: 입력 채널 32, 출력 채널 64, 커널 크기 3x3
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        # 풀링 층: 2x2 크기로 다운샘플링
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 완전 연결 층: 7x7 크기의 64개 특징 맵을 10개 클래스로 분류
        self.fc = nn.Linear(64 * 7 * 7, 10)
    
    def forward(self, x):
        # 첫 번째 합성곱 + 활성화 함수 + 풀링
        x = self.pool(F.relu(self.conv1(x)))
        # 두 번째 합성곱 + 활성화 함수 + 풀링
        x = self.pool(F.relu(self.conv2(x)))
        # 텐서 평탄화
        x = x.view(-1, 64 * 7 * 7)
        # 완전 연결 층
        x = self.fc(x)
        return x

# 모델 인스턴스 생성
model = SimpleCNN()
print(model)

이 CNN 모델은 두 개의 합성곱 층과 풀링 층, 그리고 하나의 완전 연결 층으로 구성돼 있어. 각 합성곱 층은 이미지에서 특징을 추출하고, 풀링 층은 특징 맵의 크기를 줄이는 역할을 해. 마지막으로 완전 연결 층에서 추출된 특징을 기반으로 분류를 수행하지! 🧮

간단한 비전 트랜스포머 구현하기 🔧

이번에는 PyTorch로 간단한 비전 트랜스포머 모델을 구현해볼게. 이 모델도 MNIST 데이터셋을 분류하는 데 사용할 수 있어!

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=28, patch_size=7, in_channels=1, embed_dim=64):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2
        
        # 패치 임베딩 층
        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
    
    def forward(self, x):
        # x: [batch_size, channels, height, width]
        x = self.proj(x)  # [batch_size, embed_dim, height/patch_size, width/patch_size]
        x = x.flatten(2)  # [batch_size, embed_dim, n_patches]
        x = x.transpose(1, 2)  # [batch_size, n_patches, embed_dim]
        return x

class SimpleViT(nn.Module):
    def __init__(self, img_size=28, patch_size=7, in_channels=1, embed_dim=64, num_heads=4, num_layers=2, num_classes=10):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + self.patch_embed.n_patches, embed_dim))
        
        # 트랜스포머 인코더
        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        # 분류 헤드
        self.mlp_head = nn.Linear(embed_dim, num_classes)
    
    def forward(self, x):
        # 패치 임베딩
        x = self.patch_embed(x)  # [batch_size, n_patches, embed_dim]
        
        # 클래스 토큰 추가
        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)  # [batch_size, 1+n_patches, embed_dim]
        
        # 위치 임베딩 추가
        x = x + self.pos_embed
        
        # 트랜스포머 인코더 통과
        x = x.transpose(0, 1)  # [1+n_patches, batch_size, embed_dim]
        x = self.transformer(x)
        x = x.transpose(0, 1)  # [batch_size, 1+n_patches, embed_dim]
        
        # 클래스 토큰으로 분류
        x = x[:, 0]  # [batch_size, embed_dim]
        x = self.mlp_head(x)  # [batch_size, num_classes]
        return x

# 모델 인스턴스 생성
model = SimpleViT()
print(model)

이 ViT 모델은 이미지를 패치로 나누고, 각 패치를 임베딩한 후, 트랜스포머 인코더를 통과시켜 분류를 수행해. 특히 클래스 토큰과 위치 임베딩을 추가하는 부분이 ViT의 핵심 아이디어야! 🧩

모델 학습 및 평가하기 📊

이제 구현한 모델을 학습하고 평가하는 코드를 작성해볼게. 여기서는 MNIST 데이터셋을 사용할 거야!

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 데이터 로드 및 전처리
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 모델 선택 (CNN 또는 ViT)
# model = SimpleCNN()
model = SimpleViT()

# 손실 함수 및 옵티마이저 정의
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 학습 함수
def train(model, train_loader, criterion, optimizer, device='cpu'):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
    
    train_loss = running_loss / len(train_loader)
    train_acc = 100. * correct / total
    return train_loss, train_acc

# 평가 함수
def evaluate(model, test_loader, criterion, device='cpu'):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    
    test_loss = running_loss / len(test_loader)
    test_acc = 100. * correct / total
    return test_loss, test_acc

# 학습 및 평가 실행
num_epochs = 5
for epoch in range(num_epochs):
    train_loss, train_acc = train(model, train_loader, criterion, optimizer)
    test_loss, test_acc = evaluate(model, test_loader, criterion)
    
    print(f'Epoch: {epoch+1}/{num_epochs}')
    print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')
    print(f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%')
    print('---')

이 코드를 실행하면 선택한 모델(CNN 또는 ViT)이 MNIST 데이터셋에서 학습되고 평가돼. 두 모델의 성능을 비교해보면 각각의 장단점을 직접 체험할 수 있을 거야! CNN은 빠르게 높은 정확도에 도달하고, ViT는 초기에는 느리지만 충분한 학습 후에는 비슷하거나 더 나은 성능을 보일 수 있어! 📈