[Python] pytorch와 sklearn의 train_test_split 활용하여 데이터 셋 나누기와 간단한 CNN

import torch
from torchvision import datasets
import torchvision.transforms as transforms
import os
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import torch.nn as nn
import torch.optim as optim
from torchvision.datasets import STL10
from torch.utils.data import DataLoader

#\를 그대로 사용
path2data = r'.\dataset' 

#폴더 없을 경우 생성
if not os.path.exists(path2data):
    os.mkdir(path2data)

data_transformer = transforms.Compose([transforms.ToTensor()])
train_ds = datasets.STL10(path2data, split='train', download=True, transform=data_transformer)
test0_ds = datasets.STL10(path2data, split='test', download=True, transform=data_transformer)
print(test0_ds.data.shape)
print(train_ds.data.shape)

["airplane", "bird", "car", "cat", "deer", "dog", "horse", "monkey", "ship", "truck"]

 10개의 레이블로 이루어진 STL10 데이터셋 그러나 torchvision에서 불러오게 된다면 train은 5000장 test는 8000장이였음.

validation_data, test_data = train_test_split(test0_ds, test_size=0.5, random_state=42)
# validation_data를 PyTorch 데이터셋으로 변환
validation_images = torch.stack([data[0] for data in validation_data])  # 이미지 스택
validation_labels = torch.tensor([data[1] for data in validation_data])  # 레이블 텐서

# test_data를 PyTorch 데이터셋으로 변환
test_images = torch.stack([data[0] for data in test_data])  # 이미지 스택
test_labels = torch.tensor([data[1] for data in test_data])  # 레이블 텐서

sklearn의 train_test_split을 이용하여 데이터셋을 5:5로 나누고 torch.stack() 함수를 사용하여 이미지 데이터를 하나의 텐서로 병합하고 torch.tensor() 함수를 사용하여 레이블 데이터를 텐서로 변환.

 

class CustomDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        self.data = data
        self.labels = labels
        self.transform = transform

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

    def __getitem__(self, idx):
        sample = self.data[idx]
        label = self.labels[idx]

        if self.transform:
            sample = self.transform(sample)

        return sample, label

# validation 데이터셋 생성
validation_dataset = CustomDataset(validation_images, validation_labels)

# test 데이터셋 생성
test_dataset = CustomDataset(test_images, test_labels)

# 각 데이터셋의 크기 출력
print("Validation 데이터셋 크기:", len(validation_dataset))
print("Test 데이터셋 크기:", len(test_dataset))

파이토치에서 데이터셋을 제공하는 추상 클래스를 이용하여 커스텀 데이터셋을 만들기 .

8000장이 (4000, 4000)장으로 잘 나누어진 모습.

import matplotlib.pyplot as plt
import random

# 클래스 레이블 정의 (STL10 데이터셋의 클래스 레이블)
# 여러개 랜덤 표시 해볼려고
class_labels = ["airplane", "bird", "car", "cat", "deer", "dog", "horse", "monkey", "ship", "truck"]

num_rows = 3 
num_cols = 5  
fig, axes = plt.subplots(num_rows, num_cols, figsize=(20, 12))  

for i in range(num_rows):
    for j in range(num_cols):
        index = random.randint(0, len(validation_dataset) - 1)  # 랜덤 이미지 인덱스 선택
        image, label = validation_dataset[index]
        image = image.permute(1, 2, 0) #(channels, height, width) -> (height, width, channels)
        
        axes[i, j].imshow(image)
        class_label = class_labels[label]  
        axes[i, j].set_title(f"Class: {class_label}")  
        axes[i, j].axis('off')

plt.tight_layout() 
plt.show()

시각화 할 때 채널이 젤 앞에 있기에 (channels, height, width) -> (height, width, channels)로 변경 해야함

validation 데이터셋 시각화 해보았을 때 data와 label 매칭도 잘 되었음.

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(16 * 48 * 48, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = x.view(-1, 16 * 48 * 48)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x
        
batch_size = 64
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
validation_loader = DataLoader(validation_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 모델 초기화 및 손실 함수, 최적화 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 모델 훈련
num_epochs = 5  
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}/{num_epochs}, Training Loss: {running_loss / len(train_loader)}")

    # 검증 데이터로 모델 평가
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in validation_loader:  
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    validation_accuracy = correct / total
    print(f"Epoch {epoch+1}/{num_epochs}, Validation Accuracy: {validation_accuracy * 100:.2f}%")

96 * 96 이미지가 padding=1, stride=1, kernel_size=3인 conv2d를 지나게 된다면

(W- F + 2 * P) / S  + 1 -> (96-1+1*2)/1 + 1 = 96

그후 kernel_size = 2 stride=2를 가지는 pooling을 거치면

(n * 48 * 48) 형태가 됨 

Gpu가 없어서 epoch는 5번만 Test는 언젠가 하겠지...