[Pytorch, timm] Vision Transformer로 Custom Dataset 학습하기 (이미지 분류)

[Pytorch, timm] Vision Transformer로 Custom Dataset 학습하기 (이미지 분류)

 

1. VERSION

---

• python = 3.8.10
• pytorch = 1.9.1
• timm = 0.4.12

 

2. 데이터

---

dataset 폴더 안에 각 class에 대한 폴더가 있고, class 폴더 안에 사진 파일들이 있는 상황

dataset
--class1
----class1_1.jpg
----class1_2.jpg
--class2
----class2_1.jpg
----class2_2.jpg

 

3. 코드

---

import timm
import torch
import torchvision
import torch.utils.data as data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.nn as nn
import torch.optim as optim
from torch.utils.data.dataset import random_split
from torch.utils.data import Subset

def split_dataset(dataset, train_ratio=0.8):
    # 데이터셋의 크기와 분할 지점 계산
    dataset_size = len(dataset)
    split_point = int(train_ratio * dataset_size)

    # 데이터셋의 인덱스 생성
    indices = list(range(dataset_size))

    # 학습 및 검증 세트의 인덱스 분할
    train_indices = indices[:split_point]
    val_indices = indices[split_point:]

    # Subset을 사용하여 데이터셋 분할
    train_dataset = Subset(dataset, train_indices)
    val_dataset = Subset(dataset, val_indices)

    return train_dataset, val_dataset

def train_model(model, dataloaders, criterion, optimizer, num_epochs=25):
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-' * 10)

        # 각 에포크는 학습 단계와 검증 단계를 갖습니다.
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # 모델을 학습 모드로 설정
            else:
                model.eval()   # 모델을 평가 모드로 설정

            running_loss = 0.0
            running_corrects = 0

            # 데이터를 반복
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 파라미터 경사도를 0으로 설정
                optimizer.zero_grad()

                # 순전파
                # 학습 시에만 연산 기록을 추적
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # 학습 단계인 경우 역전파 + 최적화
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # 통계
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

    return model

if __name__ == "__main__":
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # 변환(transform) 설정
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])

    # 전체 데이터셋 로드
    full_dataset = datasets.ImageFolder(root="./dataset", transform=transform)
    classes = full_dataset.classes
    num_classes = len(classes)

    # 학습 및 검증 세트 분리
    # 단순 분할 방식
    train_dataset, val_dataset = split_dataset(full_dataset, train_ratio=0.8)

    # random split 방식
    # train_size = int(0.8 * len(full_dataset))  # 학습 세트는 전체의 80%
    # val_size = len(full_dataset) - train_size  # 검증 세트는 나머지 20%
    # train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])

    dataloaders = {
        'train': data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4),
        'val': data.DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)
    }


    model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=num_classes).to(device)
    # model = timm.create_model('vit_large_patch16_224', pretrained=True, num_classes=num_classes).to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 모델 학습
    model = train_model(model, dataloaders, criterion, optimizer, num_epochs=25)

 

 

해당 코드를 실행하면 아래와 같은 Output을 얻을 수 있습니다.

Epoch 1/25
----------
train Loss: 0.2955 Acc: 0.9110
val Loss: 4.6229 Acc: 0.4961
Epoch 2/25
----------
train Loss: 0.0363 Acc: 0.9884
val Loss: 5.4199 Acc: 0.4648
Epoch 3/25
----------
train Loss: 0.0042 Acc: 0.9992
val Loss: 6.3016 Acc: 0.4945

 

이 코드는 PyTorch와 timm 라이브러리를 사용하여 이미지 분류 작업을 위한 딥러닝 모델을 학습하는 과정을 단계별로 구현한 것입니다. 코드는 크게 데이터 준비, 모델 설정, 학습 과정으로 나눌 수 있습니다. 각 단계별로 설명하겠습니다.

데이터 준비

1. 데이터셋 분할: split_dataset 함수를 통해 주어진 데이터셋을 학습 세트와 검증 세트로 분할합니다. 이 함수는 데이터셋의 크기를 기반으로 주어진 비율(train_ratio)에 따라 데이터셋을 학습과 검증용으로 나눕니다. Subset을 사용하여 원본 데이터셋에서 특정 인덱스의 데이터만 선택하여 새로운 데이터셋을 생성합니다.
2. 데이터 로딩 및 변환: PyTorch의 transforms 모듈을 사용하여 이미지에 대한 전처리 작업을 정의합니다. 이는 이미지 크기 조정, 중심 자르기, 텐서 변환, 정규화 등을 포함합니다. 이러한 변환은 ImageFolder 데이터셋을 로드할 때 적용됩니다. ImageFolder는 주어진 디렉토리 구조에서 이미지 데이터셋을 자동으로 로드하고 레이블을 지정하는 데 사용됩니다.
3. DataLoader 설정: 학습 및 검증 데이터셋을 DataLoader에 넣어 배치 처리, 셔플링, 다중 스레드를 통한 데이터 로딩을 설정합니다.

모델 설정

1. 모델 선택 및 초기화: timm 라이브러리를 사용하여 사전 학습된 vit_base_patch16_224 모델을 로드하고, 최종 분류 계층을 현재 데이터셋의 클래스 수에 맞게 조정합니다. 모델은 선택된 디바이스(CPU 또는 GPU)로 이동됩니다.
2. 손실 함수 및 옵티마이저 설정: 크로스 엔트로피 손실 함수와 SGD(확률적 경사 하강법) 옵티마이저를 사용하여 모델을 학습합니다. 옵티마이저에는 학습률과 모멘텀이 설정됩니다.

학습 과정

1. 에포크 반복: 설정된 에포크 수만큼 모델 학습 과정을 반복합니다. 각 에포크마다 학습과 검증 단계가 있습니다.
2. 모드 설정: 모델은 학습 단계에서는 train 모드로, 검증 단계에서는 eval 모드로 설정됩니다.
3. 데이터 반복 처리: 각 단계에서 DataLoader를 통해 데이터 배치를 순회하며, 모델 입력을 준비하고, 순전파를 수행한 후 손실을 계산합니다. 학습 단계에서는 역전파와 옵티마이저를 통한 파라미터 업데이트가 추가로 수행됩니다.
4. 성능 기록: 각 에포크마다 손실과 정확도를 계산하여 출력함으로써 학습과 검증 과정의 성능을 모니터링합니다.

이 코드는 딥러닝 모델 학습의 전형적인 흐름을 따르며, 데이터 준비부터 모델 학습, 성능 모니터링에 이르기까지 필요한 주요 단계를 포함합니다. timm 라이브러리의 활용은 다양한 사전 학습된 모델을 쉽게 실험할 수 있게 해주며, PyTorch 기반의 코드 구조는 모델 학습 과정의 커스터마이징을 용이하게 합니다.