[딥러닝 with Python] Vision Transformer(ViT) 파이썬 코드 구현(파이토치 / 코랩)

이번에는 지난번 알아본 Vision Transformer(ViT)의 개념을 바탕으로 파이토치에 기반한 코드를 작성해보고자합니다.

 

[딥러닝 with Python] 비전 트랜스포머(Vision Transformer / ViT) (1/2)

[딥러닝 with Python] 비전 트랜스포머(Vision Transformer / ViT) (1/2)

[딥러닝 with Python] 비전 트랜스포머(Vision Transformer / ViT) (2/2)

[딥러닝 with Python] 비전 트랜스포머(Vision Transformer / ViT) (2/2)

 

 

1. Vision Transformer(비전 트랜스포머) 코드 구현 (파이토치 / 코랩)

- 해당 코드는 파이토치를 기반으로 작성했으며, 재현성을 위해 코랩에서 작성하였습니다.

 

- 먼저, 코랩 노트북을 실행한 뒤 GPU 사용을 위해 우측 상단에 있는 부분에서 사용가능한 GPU를 선택해줍니다. 무료버전은 T4 GPU가 사용가능하며 아래 실습간에 충분한 수준입니다.

 

- 이후, "transformers" 라이브러리를 다운로드 받아줍니다. 해당 라이브러리를 huggingface의 github에 있는 라이브러리입니다.

 

https://github.com/huggingface/transformers

GitHub - huggingface/transformers: 🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

 

!pip install transformers

 

 

- 이제 학습에 활용한 ViT 모델을 가져와줍니다. 가볍게 실습하기 위해 ViT-B /16 (B는 Base(Base, Large, Huge 등등이 있지만 가장 가벼운 Base를 실습간 활용)를, 16은 이미지 패치를 16x16으로 나누었다는 의미입니다)

 * 이때, 사전에 ImageNet10K로 사전학습된 Weight를 가져와서 활용하는 전이학습을 위해 사전학습된 버전들을 가져옵니다.

from transformers import ViTModel, ViTForImageClassification, ViTFeatureExtractor

import torch

# 사전 학습된 ViTForImageClassification에서 ViTModel을 추출

feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')

backbone_model = ViTModel.from_pretrained('google/vit-base-patch16-224')  # ViTModel 사용

classifier_model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')

 

- 분류를 하기 전 ViT-B / 16 의 CLS Token이 어떤 형태인지 알아보기 위해 코드를 실행해줍니다.

# ImageNet 10K의 입력 크기에 맞는 임의의 텐서 생성

dummy_input = torch.randn(1, 3, 224, 224)

# CLS 토큰을 포함한 피처 추출

with torch.no_grad():

    outputs = backbone_model(dummy_input)  # ViTModel로 CLS 토큰 포함 히든 상태 추출

    cls_token_output = outputs.last_hidden_state[:, 0, :]  # CLS 토큰만 선택

print("CLS Token Shape:", cls_token_output.shape)

 * CLS 토큰의 shape은 [batch size, hidden state]의 형태이며, 이때 batch size는 데이터로더에따라 달라지지만, 사전에 ImagNet10K로 학습된 hiddenstate는 768차원(= 16 x 16 x 3 ) 임을 알수있습니다.

 * 이때 16 x 16 은 패치수를 의미하며 3은 채널 차원을 의미합니다. 이것을 flatten 했다는 의미가 되겠습니다.

 

- 다음은 학습에 사용할 데이터를 가져오겠습니다. 간단한 실습을 위해 CIFAR-10 Data를 활용하겠습니다.

 * 원래 CIFAR-10 데이터를 활용하게 되면 Train이 50,000개 Test가 10,000개 여서 시간이 좀 오래걸립니다.

 * 이번 실습에서는 1/20로 사이즈를 줄여서 랜덤으로 샘플링해 사용하겠으며, Train을 Train 과 Validation으로 Split 하겠습니다.

from torchvision import datasets, transforms

from torch.utils.data import DataLoader, random_split

import numpy as np

# 데이터 변환 (CIFAR-10 이미지를 ViT에 맞게 224x224 크기로 조정)

transform = transforms.Compose([

    transforms.Resize((224, 224)),   # ViT의 입력 크기에 맞게 조정

    transforms.ToTensor(),           # 이미지를 Tensor로 변환, [0, 1] 범위를 유지

])

# CIFAR-10 데이터셋 불러오기

full_train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)

test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# 전체 train 데이터셋에서 1/20만 사용하여 축소된 train 데이터셋 생성

subset_size = len(full_train_dataset) // 20

train_val_subset, _ = random_split(full_train_dataset, [subset_size, len(full_train_dataset) - subset_size])

# train 데이터셋의 80%는 train, 20%는 validation으로 분할

train_size = int(0.8 * subset_size)

validation_size = subset_size - train_size

train_subset, val_subset = random_split(train_val_subset, [train_size, validation_size])

# 테스트 데이터셋도 1/20만 사용

test_size = len(test_dataset) // 20

test_subset, _ = random_split(test_dataset, [test_size, len(test_dataset) - test_size])

# DataLoader 정의

train_loader = DataLoader(train_subset, batch_size=128, shuffle=True)

val_loader = DataLoader(val_subset, batch_size=128, shuffle=False)

test_loader = DataLoader(test_subset, batch_size=128, shuffle=False)

print("Training data samples (subset):", len(train_loader.dataset))

print("Validation data samples (subset):", len(val_loader.dataset))

print("Test data samples (subset):", len(test_loader.dataset))

 

- 다음은 사전학습된 Weight가 새로운 데이터 분류에 용이한지 알아보기 위해 Linear Probe를 진행해보겠습니다. Linear Probe는 사전학습된 모델의 마지막 부분에 Down Stream Task를 위한 1층의 Linear Layer를 추가하는 것으로 자세한 개념은 아래 포스팅을 참조해주시기 바랍니다. 이때 사전학습된 Weight는 고정(Freeze)시킵니다.

[개념 정리] Linear probing이란?

[개념 정리] Linear probing이란?

 

- Pretrained 된 ViT에 1층의 Linear Layer를 통해 분류한 Linear Probing의 결과는 아래와 같습니다.

import torch

import torch.nn as nn

from transformers import ViTModel

from torch.optim import Adam

from torch.cuda.amp import autocast, GradScaler

# 모델 준비

backbone_model = ViTModel.from_pretrained('google/vit-base-patch16-224')

for param in backbone_model.parameters():

    param.requires_grad = False  # ViTModel 가중치 고정

# Linear 분류기 정의

classifier = nn.Linear(backbone_model.config.hidden_size, 10)  # CIFAR-10은 10개의 클래스

# GPU 설정

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

backbone_model.to(device)

classifier.to(device)

# 손실 함수와 옵티마이저 정의

criterion = nn.CrossEntropyLoss()

optimizer = Adam(classifier.parameters(), lr=1e-3)

# 혼합 정밀도 학습을 위한 스케일러

scaler = GradScaler()

# 학습 함수

def train(backbone_model, classifier, dataloader, optimizer, criterion, device, scaler):

    backbone_model.eval()  # Backbone 모델은 고정 (학습하지 않음)

    classifier.train()

    total_loss = 0

    for batch in dataloader:

        inputs, labels = batch

        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        with torch.no_grad():

            with autocast():  # 혼합 정밀도

                outputs = backbone_model(inputs).last_hidden_state[:, 0, :]  # CLS 토큰만 추출

        with autocast():

            logits = classifier(outputs)

            loss = criterion(logits, labels)

        scaler.scale(loss).backward()

        scaler.step(optimizer)

        scaler.update()

        total_loss += loss.item()

    return total_loss / len(dataloader)

# 평가 함수

def evaluate(backbone_model, classifier, dataloader, criterion, device):

    backbone_model.eval()

    classifier.eval()

    total_loss = 0

    correct = 0

    total = 0

    with torch.no_grad():

        for batch in dataloader:

            inputs, labels = batch

            inputs, labels = inputs.to(device), labels.to(device)

            outputs = backbone_model(inputs).last_hidden_state[:, 0, :]

            logits = classifier(outputs)

            loss = criterion(logits, labels)

            total_loss += loss.item()

            _, predicted = torch.max(logits, 1)

            total += labels.size(0)

            correct += (predicted == labels).sum().item()

    accuracy = correct / total

    avg_loss = total_loss / len(dataloader)

    return accuracy, avg_loss

# 학습 및 평가 + Early Stopping

num_epochs = 30

patience = 5  # Early stopping patience

best_val_loss = float("inf")

early_stopping_counter = 0

for epoch in range(num_epochs):

    train_loss = train(backbone_model, classifier, train_loader, optimizer, criterion, device, scaler)

    val_accuracy, val_loss = evaluate(backbone_model, classifier, val_loader, criterion, device)

    print(f"Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}")

    # Early stopping 체크

    if val_loss < best_val_loss:

        best_val_loss = val_loss

        early_stopping_counter = 0  # 개선된 경우 카운터 초기화

        best_model_wts = classifier.state_dict()  # 가장 좋은 모델 가중치 저장

    else:

        early_stopping_counter += 1

        if early_stopping_counter >= patience:

            print("Early stopping triggered")

            break

# 최적 모델로 복원

classifier.load_state_dict(best_model_wts)

# 최종 테스트 정확도 평가

test_accuracy, test_loss = evaluate(backbone_model, classifier, test_loader, criterion, device)

print(f"Test Accuracy: {test_accuracy:.4f}, Test Loss: {test_loss:.4f}")

 

* 88.8%의 테스트 데이터셋에 대한 정확도를 보여주었습니다.

 

- 이번에는 Pretrained 되어있는 Weight를 고정하지 않고 학습가능하게 해서 분류를 진행해보겠습니다. 이번에도 마찬가지로 마지 부분에 1개의 Layer만을 활용했습니다.

 

import torch

import torch.nn as nn

from transformers import ViTModel

from torch.optim import Adam

from torch.cuda.amp import autocast, GradScaler

# Fine-tuning 전에 모델을 초기화하여 이전 학습 영향 방지

fine_tune_model = ViTModel.from_pretrained('google/vit-base-patch16-224')

for param in fine_tune_model.parameters():

    param.requires_grad = True  # 모든 레이어를 학습 가능하게 설정

# Fine-tuning용 Linear 분류기 정의

fine_tune_classifier = nn.Linear(fine_tune_model.config.hidden_size, 10)  # CIFAR-10 클래스 수

# GPU 설정

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

fine_tune_model.to(device)

fine_tune_classifier.to(device)

# 손실 함수와 옵티마이저 정의

criterion = nn.CrossEntropyLoss()

optimizer = Adam(list(fine_tune_model.parameters()) + list(fine_tune_classifier.parameters()), lr=1e-3)

# 혼합 정밀도 학습을 위한 스케일러

scaler = GradScaler()

# 학습 함수

def train(model, classifier, dataloader, optimizer, criterion, device, scaler):

    model.train()

    classifier.train()

    total_loss = 0

    for batch in dataloader:

        inputs, labels = batch

        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        with autocast():  # 혼합 정밀도

            outputs = model(inputs).last_hidden_state[:, 0, :]  # CLS 토큰만 추출

            logits = classifier(outputs)

            loss = criterion(logits, labels)

        scaler.scale(loss).backward()

        scaler.step(optimizer)

        scaler.update()

        total_loss += loss.item()

    return total_loss / len(dataloader)

# 평가 함수

def evaluate(model, classifier, dataloader, criterion, device):

    model.eval()

    classifier.eval()

    total_loss = 0

    correct = 0

    total = 0

    with torch.no_grad():

        for batch in dataloader:

            inputs, labels = batch

            inputs, labels = inputs.to(device), labels.to(device)

            outputs = model(inputs).last_hidden_state[:, 0, :]

            logits = classifier(outputs)

            loss = criterion(logits, labels)

            total_loss += loss.item()

            _, predicted = torch.max(logits, 1)

            total += labels.size(0)

            correct += (predicted == labels).sum().item()

    accuracy = correct / total

    avg_loss = total_loss / len(dataloader)

    return accuracy, avg_loss

# 학습 및 평가

num_epochs = 30

best_val_loss = float("inf")

early_stopping_counter = 0

patience = 5

for epoch in range(num_epochs):

    train_loss = train(fine_tune_model, fine_tune_classifier, train_loader, optimizer, criterion, device, scaler)

    val_accuracy, val_loss = evaluate(fine_tune_model, fine_tune_classifier, val_loader, criterion, device)

    print(f"Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}")

    # Early stopping 체크

    if val_loss < best_val_loss:

        best_val_loss = val_loss

        early_stopping_counter = 0  # 개선된 경우 카운터 초기화

        best_model_wts = fine_tune_classifier.state_dict()  # 가장 좋은 모델 가중치 저장

    else:

        early_stopping_counter += 1

        if early_stopping_counter >= patience:

            print("Early stopping triggered")

            break

# 최적 모델로 복원

fine_tune_classifier.load_state_dict(best_model_wts)

# 최종 테스트 정확도 평가

test_accuracy, test_loss = evaluate(fine_tune_model, fine_tune_classifier, test_loader, criterion, device)

print(f"Test Accuracy: {test_accuracy:.4f}, Test Loss: {test_loss:.4f}")

 

 

 * Test 정확도는 40.8%가 나왔습니다. 아무래도 새로운 데이터이고 데이터의 숫자도 적다보니 Fine tuning이 충분히 되지 않아서 더 안좋은 결과가 나온 것 같습니다.

 

- 이번에는 Pretrained Weight는 고정하고, 다른 기법들을 활용해서 성능을 향상시켜보겠습니다.

* 성능향상에 활용한 기법은 Data Augmentation(Random Horizontal Flip, Random Rotation, Color Jittering)과 2층의 MLP Classifier(ReLU,Drop Out 포함)를 활용했습니다.

import torch

import torch.nn as nn

from torch.optim import Adam, lr_scheduler

from torch.cuda.amp import autocast, GradScaler

from torchvision import transforms

# 데이터 증강 기법 추가

train_transform = transforms.Compose([

    transforms.Resize((224, 224)),

    transforms.RandomHorizontalFlip(),   # 랜덤 좌우 반전

    transforms.RandomRotation(10),       # 랜덤 회전

    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),

    transforms.ToTensor(),

])

# 모델과 분류기 정의 (기존 모델 초기화)

enhanced_model = ViTModel.from_pretrained('google/vit-base-patch16-224')

for param in enhanced_model.parameters():

    param.requires_grad = True  # 모든 레이어를 학습 가능하게 설정

# Fine-tuning을 위한 복잡한 MLP 분류기

enhanced_classifier = nn.Sequential(

    nn.Linear(enhanced_model.config.hidden_size, 512),  # 중간 레이어 추가

    nn.ReLU(),

    nn.Dropout(0.3),

    nn.Linear(512, 10)  # CIFAR-10 클래스 수

)

# GPU 설정

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

enhanced_model.to(device)

enhanced_classifier.to(device)

# 손실 함수와 옵티마이저, 스케줄러 정의

criterion = nn.CrossEntropyLoss()

optimizer = Adam(list(enhanced_model.parameters()) + list(enhanced_classifier.parameters()), lr=1e-3)

scheduler = lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.7)  # 3 에포크마다 학습률 감소

# 혼합 정밀도 학습을 위한 스케일러

scaler = GradScaler()

# 학습 함수

def train(model, classifier, dataloader, optimizer, criterion, device, scaler):

    model.train()

    classifier.train()

    total_loss = 0

    for batch in dataloader:

        inputs, labels = batch

        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        with autocast():  # 혼합 정밀도

            outputs = model(inputs).last_hidden_state[:, 0, :]  # CLS 토큰만 추출

            logits = classifier(outputs)

            loss = criterion(logits, labels)

        scaler.scale(loss).backward()

        scaler.step(optimizer)

        scaler.update()

        total_loss += loss.item()

    return total_loss / len(dataloader)

# 평가 함수

def evaluate(model, classifier, dataloader, criterion, device):

    model.eval()

    classifier.eval()

    total_loss = 0

    correct = 0

    total = 0

    with torch.no_grad():

        for batch in dataloader:

            inputs, labels = batch

            inputs, labels = inputs.to(device), labels.to(device)

            outputs = model(inputs).last_hidden_state[:, 0, :]

            logits = classifier(outputs)

            loss = criterion(logits, labels)

            total_loss += loss.item()

            _, predicted = torch.max(logits, 1)

            total += labels.size(0)

            correct += (predicted == labels).sum().item()

    accuracy = correct / total

    avg_loss = total_loss / len(dataloader)

    return accuracy, avg_loss

# 학습 및 평가 루프 + Early Stopping

num_epochs = 50

best_val_loss = float("inf")

early_stopping_counter = 0

patience = 5

for epoch in range(num_epochs):

    train_loss = train(enhanced_model, enhanced_classifier, train_loader, optimizer, criterion, device, scaler)

    val_accuracy, val_loss = evaluate(enhanced_model, enhanced_classifier, val_loader, criterion, device)

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

    print(f"Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}")

    # Early stopping 체크

    if val_loss < best_val_loss:

        best_val_loss = val_loss

        early_stopping_counter = 0

        best_model_wts = enhanced_classifier.state_dict()

    else:

        early_stopping_counter += 1

        if early_stopping_counter >= patience:

            print("Early stopping triggered")

            break

# 최적 모델 복원

enhanced_classifier.load_state_dict(best_model_wts)

# 테스트 평가

test_accuracy, test_loss = evaluate(enhanced_model, enhanced_classifier, test_loader, criterion, device)

print(f"Test Accuracy: {test_accuracy:.4f}, Test Loss: {test_loss:.4f}")

 

* 결과는 위와 같이 Test Accuracy가 향상됨을 확인할 수 있습니다.