[CDAL Research] CNN Model for Image Classification

 

 

---

1. Module Import & GPU Usage

 

 

• Basic Module import

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torchvision.transforms import autoaugment
from torch.utils.data import DataLoader, random_split
import numpy as np
import os
from PIL import Image
import matplotlib.pyplot as plt

 

• M1 Macbook에서 GPU 사용: mps 이용

device = torch.device("mps:0" if torch.backends.mps.is_available() else "cpu")
print(f"Using {device} device")

 

• GPU 사용 전 점검사항

print (f"PyTorch version:{torch.__version__}") # 1.12.1 이상
print(f"MPS 장치를 지원하도록 build 되었는지: {torch.backends.mps.is_built()}") # True 여야 합니다.
print(f"MPS 장치가 사용 가능한지: {torch.backends.mps.is_available()}") # True 여야 합니다.
!python -c 'import platform;print(platform.platform())'

PyTorch version:2.1.2
MPS 장치를 지원하도록 build 되었는지: True
MPS 장치가 사용 가능한지: True
macOS-14.2-arm64-arm-64bit

 

 

 

---

2. Dataset & Basic Image Preprocessing

 

 

 

 

• Dataset 준비하기
  
  • Datset 생성 및 합치기
  • Train / Validation / Test Dataset 분할
  • DataLoader 생성

 

from torch.utils.data import Dataset, DataLoader, ConcatDataset, random_split
from PIL import Image
import os
import torchvision.transforms as transforms


class CustomImageDataset(Dataset):
    def __init__(self, img_dir, transform=None, label=0):
        self.img_dir = img_dir
        self.transform = transform
        self.label = label
        self.img_paths = [
            os.path.join(img_dir, img_file)
            for img_file in os.listdir(img_dir)
            if img_file.lower().endswith(("png", "jpg", "jpeg"))
        ]

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

    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, self.label


# 이미지 경로
root_dir = "/Users/eric/Documents/Research/CDAL/2023_Winter/summer2winter_yosemite"

## 이미지 변환 설정
transform = transforms.Compose(
    [
        transforms.Resize((256, 256)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

# 데이터셋 생성
train_dataset_A = CustomImageDataset(
    os.path.join(root_dir, "trainA"), transform=transform, label=0
)
train_dataset_B = CustomImageDataset(
    os.path.join(root_dir, "trainB"), transform=transform, label=1
)

test_dataset_A = CustomImageDataset(
    os.path.join(root_dir, "testA"), transform=transform, label=0
)
test_dataset_B = CustomImageDataset(
    os.path.join(root_dir, "testB"), transform=transform, label=1
)

# 데이터셋 합치기
train_dataset = ConcatDataset([train_dataset_A, train_dataset_B])
test_dataset = ConcatDataset([test_dataset_A, test_dataset_B])

# 검증 데이터셋 분할
val_split = 0.2
train_size = int((1 - val_split) * len(train_dataset))
val_size = len(train_dataset) - train_size
train_dataset_split, val_dataset_split = random_split(
    train_dataset, [train_size, val_size]
)

# DataLoader 생성
train_loader = DataLoader(train_dataset_split, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset_split, batch_size=32, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

 

 

• 기본 이미지 전처리 & Augmentation

# Data Preprocessing & Augmentation 설정
transform = transforms.Compose(
    [
        transforms.Resize((256, 256)),
        transforms.RandomHorizontalFlip(),
        autoaugment.RandAugment(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

 

 

 

---

3. CNN Modeling

 

 

 

• Class Weight 계산: CE Loss Function에 적용

# 클래스 가중치 계산
class_counts = [0, 0]  # 클래스 수에 따라 조정
for _, label in train_dataset:
    class_counts[label] += 1

class_weights = [sum(class_counts) / c for c in class_counts]
class_weights = torch.tensor(class_weights, dtype=torch.float).to(device)

 

 

• CNNModel 정의

class CNNModel(nn.Module):
    def __init__(self):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(256 * 32 * 32, 512)
        self.fc2 = nn.Linear(512, 2)
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 256 * 32 * 32)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)

        return x


# Model(GPU로 이동), Loss Function, Optimizer 설정
model = CNNModel().to(device)
# 손실 함수에 클래스 가중치 적용
criterion = nn.CrossEntropyLoss(weight=class_weights)
optimizer = optim.Adam(model.parameters(), lr=0.001)

 

 

 

 

---

4. Mixup / CutMix Preprocessing 적용

 

 

 

 

 

Mixup Method

 

 

CutMix Method

 

 

 

 

• Additional Preprocessing
  • CutMix
  • Mixup

 

# Mixup/CutMix 데이터 처리
def mixup_data(x, y, alpha=1.0):
    """Mixup 데이터 처리 구현"""
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1
    batch_size = x.size(0)
    index = torch.randperm(batch_size).to(device)
    mixed_x = lam * x + (1 - lam) * x[index, :]
    y_a, y_b = y, y[index]
    return mixed_x, y_a, y_b, lam


def cutmix_data(x, y, alpha=1.0):
    """CutMix 데이터 처리 구현"""
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1
    batch_size = x.size(0)
    index = torch.randperm(batch_size).to(device)

    # 랜덤한 위치와 크기 결정
    W, H = x.size(2), x.size(3)
    cut_rat = np.sqrt(1.0 - lam)
    cut_w = int(W * cut_rat)
    cut_h = int(H * cut_rat)

    # 유니폼 분포에서 중심 위치 결정
    cx = np.random.randint(W)
    cy = np.random.randint(H)

    # 패치 영역 결정
    bbx1 = np.clip(cx - cut_w // 2, 0, W)
    bby1 = np.clip(cy - cut_h // 2, 0, H)
    bbx2 = np.clip(cx + cut_w // 2, 0, W)
    bby2 = np.clip(cy + cut_h // 2, 0, H)

    # 원본 이미지에 패치 적용
    x[:, :, bbx1:bbx2, bby1:bby2] = x[index, :, bbx1:bbx2, bby1:bby2]
    lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (W * H))

    y_a, y_b = y, y[index]
    return x, y_a, y_b, lam


# Mixup/CutMix 손실 함수
def mixup_criterion(criterion, pred, y_a, y_b, lam):
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)

 

 

 

---

5. Train, Validate, Test 함수 정의

 

 

# 훈련 및 검증 함수 정의
def train(model, train_loader, optimizer, criterion, alpha=1.0):
    model.train()
    train_loss = 0.0
    correct = 0
    total = 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        images, labels_a, labels_b, lam = mixup_data(images, labels, alpha)
        optimizer.zero_grad()
        outputs = model(images)
        loss = mixup_criterion(criterion, outputs, labels_a, labels_b, lam)
        train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        loss.backward()
        optimizer.step()
    train_loss /= total
    train_accuracy = 100 * correct / total
    return train_loss, train_accuracy


def validate(model, val_loader, criterion):
    model.eval()
    val_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            val_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    val_loss /= total
    val_accuracy = 100 * correct / total
    return val_loss, val_accuracy


def test(model, test_loader):
    model.eval()
    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)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100 * correct / total

 

 

 

---

6. Train & Validation & Test Implementation

 

 

# 손실 및 정확도 추적을 위한 리스트
train_losses, train_accuracies, val_losses, val_accuracies = [], [], [], []

# 훈련 및 검증
epochs = 25
for epoch in range(epochs):
    train_loss, train_accuracy = train(model, train_loader, optimizer, criterion)
    val_loss, val_accuracy = validate(model, val_loader, criterion)
    train_losses.append(train_loss)
    train_accuracies.append(train_accuracy)
    val_losses.append(val_loss)
    val_accuracies.append(val_accuracy)
    print(f"Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.2f}%, Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.2f}%")

# 테스트 세트에서의 성능
test_accuracy = test(model, test_loader)
print(f"Test Accuracy: {test_accuracy:.2f}%")

poch 1/40, Train Loss: 0.0441, Train Accuracy: 55.70%, Val Loss: 0.0215, Val Accuracy: 53.08%
Epoch 2/40, Train Loss: 0.0210, Train Accuracy: 58.38%, Val Loss: 0.0212, Val Accuracy: 67.88%
Epoch 3/40, Train Loss: 0.0204, Train Accuracy: 60.78%, Val Loss: 0.0180, Val Accuracy: 71.98%
Epoch 4/40, Train Loss: 0.0199, Train Accuracy: 60.32%, Val Loss: 0.0173, Val Accuracy: 71.75%
Epoch 5/40, Train Loss: 0.0193, Train Accuracy: 63.11%, Val Loss: 0.0174, Val Accuracy: 73.80%
Epoch 6/40, Train Loss: 0.0189, Train Accuracy: 63.97%, Val Loss: 0.0178, Val Accuracy: 72.44%
Epoch 7/40, Train Loss: 0.0184, Train Accuracy: 66.08%, Val Loss: 0.0166, Val Accuracy: 73.58%
Epoch 8/40, Train Loss: 0.0179, Train Accuracy: 64.82%, Val Loss: 0.0159, Val Accuracy: 75.17%
Epoch 9/40, Train Loss: 0.0182, Train Accuracy: 66.13%, Val Loss: 0.0157, Val Accuracy: 77.45%
Epoch 10/40, Train Loss: 0.0183, Train Accuracy: 65.17%, Val Loss: 0.0153, Val Accuracy: 76.54%
Epoch 11/40, Train Loss: 0.0180, Train Accuracy: 65.22%, Val Loss: 0.0153, Val Accuracy: 77.22%
Epoch 12/40, Train Loss: 0.0176, Train Accuracy: 68.24%, Val Loss: 0.0176, Val Accuracy: 74.94%
Epoch 13/40, Train Loss: 0.0174, Train Accuracy: 64.65%, Val Loss: 0.0158, Val Accuracy: 77.45%
Epoch 14/40, Train Loss: 0.0175, Train Accuracy: 70.70%, Val Loss: 0.0161, Val Accuracy: 76.31%
Epoch 15/40, Train Loss: 0.0163, Train Accuracy: 68.93%, Val Loss: 0.0162, Val Accuracy: 77.22%
Epoch 16/40, Train Loss: 0.0161, Train Accuracy: 68.36%, Val Loss: 0.0162, Val Accuracy: 77.68%
Epoch 17/40, Train Loss: 0.0162, Train Accuracy: 68.81%, Val Loss: 0.0164, Val Accuracy: 77.22%
Epoch 18/40, Train Loss: 0.0152, Train Accuracy: 72.86%, Val Loss: 0.0165, Val Accuracy: 77.22%
Epoch 19/40, Train Loss: 0.0147, Train Accuracy: 69.50%, Val Loss: 0.0159, Val Accuracy: 77.45%
Epoch 20/40, Train Loss: 0.0143, Train Accuracy: 74.46%, Val Loss: 0.0160, Val Accuracy: 77.45%
Epoch 21/40, Train Loss: 0.0135, Train Accuracy: 71.32%, Val Loss: 0.0171, Val Accuracy: 79.95%
Epoch 22/40, Train Loss: 0.0145, Train Accuracy: 72.18%, Val Loss: 0.0168, Val Accuracy: 79.50%
Epoch 23/40, Train Loss: 0.0120, Train Accuracy: 71.04%, Val Loss: 0.0170, Val Accuracy: 80.18%
Epoch 24/40, Train Loss: 0.0133, Train Accuracy: 72.63%, Val Loss: 0.0178, Val Accuracy: 79.27%
Epoch 25/40, Train Loss: 0.0137, Train Accuracy: 71.32%, Val Loss: 0.0171, Val Accuracy: 79.27%
Epoch 26/40, Train Loss: 0.0125, Train Accuracy: 74.29%, Val Loss: 0.0165, Val Accuracy: 78.36%
Epoch 27/40, Train Loss: 0.0118, Train Accuracy: 72.46%, Val Loss: 0.0190, Val Accuracy: 77.68%
Epoch 28/40, Train Loss: 0.0127, Train Accuracy: 72.29%, Val Loss: 0.0187, Val Accuracy: 78.59%
Epoch 29/40, Train Loss: 0.0124, Train Accuracy: 74.86%, Val Loss: 0.0168, Val Accuracy: 77.90%
Epoch 30/40, Train Loss: 0.0128, Train Accuracy: 72.86%, Val Loss: 0.0191, Val Accuracy: 78.59%
Epoch 31/40, Train Loss: 0.0126, Train Accuracy: 72.41%, Val Loss: 0.0173, Val Accuracy: 78.36%
Epoch 32/40, Train Loss: 0.0115, Train Accuracy: 76.23%, Val Loss: 0.0174, Val Accuracy: 77.90%
Epoch 33/40, Train Loss: 0.0123, Train Accuracy: 74.63%, Val Loss: 0.0168, Val Accuracy: 80.41%
Epoch 34/40, Train Loss: 0.0125, Train Accuracy: 78.05%, Val Loss: 0.0165, Val Accuracy: 78.36%
Epoch 35/40, Train Loss: 0.0128, Train Accuracy: 76.00%, Val Loss: 0.0178, Val Accuracy: 76.77%
Epoch 36/40, Train Loss: 0.0127, Train Accuracy: 72.46%, Val Loss: 0.0173, Val Accuracy: 79.50%
Epoch 37/40, Train Loss: 0.0127, Train Accuracy: 74.34%, Val Loss: 0.0178, Val Accuracy: 78.13%
Epoch 38/40, Train Loss: 0.0117, Train Accuracy: 75.66%, Val Loss: 0.0192, Val Accuracy: 80.18%
Epoch 39/40, Train Loss: 0.0119, Train Accuracy: 69.61%, Val Loss: 0.0177, Val Accuracy: 79.27%
Epoch 40/40, Train Loss: 0.0118, Train Accuracy: 80.96%, Val Loss: 0.0164, Val Accuracy: 79.27%
Test Accuracy: 75.14%

 

 

• Graph
  • Loss b/w train & validation dataset
  • Accuracy b/w train & validation dataset

 

# 손실 및 정확도 시각화
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs + 1), train_losses, label="Training Loss")
plt.plot(range(1, epochs + 1), val_losses, label="Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training and Validation Loss")
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(range(1, epochs + 1), train_accuracies, label="Training Accuracy")
plt.plot(range(1, epochs + 1), val_accuracies, label="Validation Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy (%)")
plt.title("Training and Validation Accuracy")
plt.legend()

plt.tight_layout()
plt.show()