[Adapter]

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import Subset, DataLoader
from torch.utils.data.dataset import Dataset
from PIL import Image
import timm

# ViT 모델 로드
def load_vit_model():
    model = timm.create_model('vit_base_patch16_224', pretrained=True)
    preprocess = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    ])
    return model, preprocess

# Adapter 모듈 정의
class Adapter(nn.Module):
    def __init__(self, in_features, bottleneck_size=64):
        super(Adapter, self).__init__()
        self.down_project = nn.Linear(in_features, bottleneck_size)
        self.relu = nn.ReLU()
        self.up_project = nn.Linear(bottleneck_size, in_features)

    def forward(self, x):
        x = self.down_project(x)
        x = self.relu(x)
        x = self.up_project(x)
        return x

# Filtered CIFAR-10 데이터셋 정의
class FilteredCIFAR10(Dataset):
    def __init__(self, root, train=True, transform=None, download=False):
        self.cifar10 = datasets.CIFAR10(root=root, train=train, transform=transform, download=download)
        self.data = []
        self.targets = []
        for img, target in zip(self.cifar10.data, self.cifar10.targets):
            if target < 8:  # Only keep classes 0-7
                self.data.append(img)
                self.targets.append(target)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img, target = self.data[idx], self.targets[idx]
        img = Image.fromarray(img)
        if self.cifar10.transform:
            img = self.cifar10.transform(img)
        return img, target

# ViT 미세 조정 모델 정의 (Adapter 포함)
class ViTFineTuningWithAdapter(nn.Module):
    def __init__(self, base_model, num_classes):
        super(ViTFineTuningWithAdapter, self).__init__()
        self.base_model = base_model
        self.adapters = nn.ModuleList([Adapter(base_model.blocks[i].mlp.fc1.in_features) for i in range(len(base_model.blocks))])
        self.fc = nn.Linear(self.base_model.head.in_features, num_classes)
        self.base_model.head = self.fc  # 모델의 헤드 부분을 새로운 FC 레이어로 교체

    def forward(self, x):
        x = self.base_model.patch_embed(x)
        cls_tokens = self.base_model.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.base_model.pos_drop(x + self.base_model.pos_embed)
        
        for i, blk in enumerate(self.base_model.blocks):
            x = blk(x) + self.adapters[i](x)
        
        x = self.base_model.norm(x)
        x = x[:, 0]  # 첫 번째 클래스 토큰만 사용
        x = self.fc(x)
        return x

# GPU 사용 설정
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ViT 모델과 전처리 함수 로드
base_model, preprocess = load_vit_model()
model = ViTFineTuningWithAdapter(base_model, num_classes=8).to(device)  # num_classes 변경

# 기본 모델의 모든 파라미터 고정
for name, param in model.base_model.named_parameters():
    param.requires_grad = False

# Adapter와 FC 레이어만 학습 가능하게 설정
for name, param in model.named_parameters():
    if 'adapters' in name or 'fc' in name:
        param.requires_grad = True

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)  # 학습률 스케줄러 추가

# 전처리 함수로 데이터 로딩
train_dataset = FilteredCIFAR10(root='./data', train=True, download=True, transform=preprocess)
test_dataset = FilteredCIFAR10(root='./data', train=False, download=True, transform=preprocess)

# 훈련 및 테스트 데이터셋에서 100개의 샘플 선택
train_subset = Subset(train_dataset, range(100))
test_subset = Subset(test_dataset, range(100))

# 데이터 로더
train_loader = DataLoader(train_subset, batch_size=64, shuffle=True, num_workers=2)
test_loader = DataLoader(test_subset, batch_size=64, shuffle=False, num_workers=2)

# 훈련 루프 (간단히)
num_epochs = 20

for epoch in range(num_epochs):
    model.train()
    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()

    scheduler.step()  # 학습률 스케줄러 스텝

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

    # 에포크마다 정확도 계산
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        accuracy = 100 * correct / total
        print(f'Accuracy after epoch {epoch+1}: {accuracy:.2f} %')

print('Finished Training')

# 최종 정확도
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print(f'Final Test Accuracy: {100 * correct / total:.2f} %')

수정한 코드

import torch
import torch.nn as nn
import timm
import torchvision.transforms as transforms

# ViT 모델 로드
def load_vit_model():
    model = timm.create_model('vit_base_patch16_224', pretrained=True)
    preprocess = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    ])
    return model, preprocess

# Adapter 모듈 정의 (논문 기반)
class Adapter(nn.Module):
    def __init__(self, in_features, bottleneck_size=64):
        super(Adapter, self).__init__()
        self.down_project = nn.Linear(in_features, bottleneck_size)
        self.relu = nn.ReLU()
        self.up_project = nn.Linear(bottleneck_size, in_features)
        self.layernorm = nn.LayerNorm(in_features)  # 논문에서는 LayerNorm을 사용하기도 함

    def forward(self, x):
        # 원본 특징을 보존하기 위해 잔차 연결 (Residual Connection) 추가
        residual = x
        x = self.layernorm(x)
        x = self.down_project(x)  # 차원 축소
        x = self.relu(x)           # 비선형성 추가
        x = self.up_project(x)     # 차원 복원
        return x + residual         # 잔차 연결

# ViT 미세 조정 모델 정의 (Adapter 포함)
class ViTFineTuningWithAdapter(nn.Module):
    def __init__(self, base_model, num_classes):
        super(ViTFineTuningWithAdapter, self).__init__()
        self.base_model = base_model
        # 각 Transformer 블록에 Adapter 삽입
        self.adapters = nn.ModuleList([
            Adapter(base_model.blocks[i].mlp.fc1.in_features) for i in range(len(base_model.blocks))
        ])
        self.fc = nn.Linear(self.base_model.head.in_features, num_classes)
        self.base_model.head = self.fc  # 모델의 헤드 부분을 새로운 FC 레이어로 교체

    def forward(self, x):
        x = self.base_model.patch_embed(x)
        cls_tokens = self.base_model.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.base_model.pos_drop(x + self.base_model.pos_embed)
        
        for i, blk in enumerate(self.base_model.blocks):
            x = blk(x)
            x = self.adapters[i](x)  # 각 Transformer 블록 뒤에 Adapter 적용
        
        x = self.base_model.norm(x)
        x = x[:, 0]  # 첫 번째 클래스 토큰만 사용
        x = self.fc(x)
        return x

# 나머지 코드 (데이터 로딩, 훈련 루프 등)는 동일하게 유지