[PyTorch] 이미지 Classification - CNN (1) CIFAR10 dataset

1. 데이터 준비하기

torchvision에 built-in 데이터셋인 CIFAR10 dataset에 CNN모델을 이용한 이미지 classification을 진행해보자.

 

CIFAR10 dataset은 32x32사이즈, (비행기, 차, 새, 고양이, 사슴, 개, 개구리, 말, 배, 트럭)으로 총 10가지 카테고리로 분류된다. 각 카테고리별 6천개, 총 6만개의 이미지이고 training image 5만개, test image 1만개로 구성되어 있다. 

from torchvision.datasets import CIFAR10
from torchvision.transforms import transforms
from torch.utils.data import DataLoader

아래 링크에서 CIFAR10이외에 다양한 이미지 관련 빌트인 데이터 셋을 확인할 수 있다.

https://pytorch.org/vision/stable/datasets.html

 

torchvision에서 tranforms 모듈을 사용하면 이미지 데이터 전처리를 쉽게 할 수 있다. tranforms를 이용해 일괄적으로 이미지 size 변환, tensor 변환, normalize 등을 처리하자.

# Define transformations for the training and test sets
transformations = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train dataset과 test dataset을 각각 불러오자. 이때 위에서 정의한 transformation을 사용하면 내가 정한 포맷대로 데이터를 불러올 수 있다.

# Create an instance for training. 
# When we run this code for the first time, the CIFAR10 train dataset will be downloaded locally. 
train_set =CIFAR10(root="./data",train=True,transform=transformations,download=True)

# Create an instance for testing, note that train is set to False.
# When we run this code for the first time, the CIFAR10 test dataset will be downloaded locally. 
test_set = CIFAR10(root="./data", train=False, transform=transformations, download=True)

classes = train_set.classes

 

pytorch의 DataLoder는 train / test에서 데이터를 편리하게 불러오기 위해 사용하는 모듈이다.

 

batch size, shuffle 여부 등을 지정하면 알아서 처리해주기 때문에 데이터를 배치 사이즈에 맞춰서 자르고, 랜덤하게 섞고..등 이런 번거로운 코드를 생략할 수 있다. 

# CIFAR10 dataset consists of 50K training images. We define the batch size of 10 to load 5,000 batches of images.
batch_size = 10
number_of_labels = 10

# Create a loader for the training set which will read the data within batch size and put into memory.
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0)
print("The number of images in a training set is: ", len(train_loader)*batch_size)

# Create a loader for the test set which will read the data within batch size and put into memory. 
# Note that each shuffle is set to false for the test loader.
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=0)
print("The number of images in a test set is: ", len(test_loader)*batch_size)
print("The number of batches per epoch is: ", len(train_loader))

 

The number of images in a training set is:  50000
The number of images in a test set is:  10000
The number of batches per epoch is:  5000

 

2. model

총 14개의 layer로 구성된 CNN 모델이다.

 

Input -> Conv -> BatchNorm -> ReLU 

-> Conv -> BatchNorm -> ReLU 

-> MaxPool 

-> Conv -> BatchNorm -> ReLU 

-> Conv -> BatchNorm -> ReLU 

-> Linear -> output

# Define a convolution neural network
class Network(nn.Module):
    '''
    14 layer
    '''
    def __init__(self):
        super(Network, self).__init__()
        
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=5, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(12)
        self.conv2 = nn.Conv2d(in_channels=12, out_channels=12, kernel_size=5, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(12)
        self.pool = nn.MaxPool2d(2,2)
        self.conv4 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=5, stride=1, padding=1)
        self.bn4 = nn.BatchNorm2d(24)
        self.conv5 = nn.Conv2d(in_channels=24, out_channels=24, kernel_size=5, stride=1, padding=1)
        self.bn5 = nn.BatchNorm2d(24)
        self.fc1 = nn.Linear(24*10*10, 10)
        
    def forward(self, input):
        output = F.relu(self.bn1(self.conv1(input)))      
        output = F.relu(self.bn2(self.conv2(output)))     
        output = self.pool(output)                        
        output = F.relu(self.bn4(self.conv4(output)))     
        output = F.relu(self.bn5(self.conv5(output)))     
        output = output.view(-1, 24*10*10)
        output = self.fc1(output)

        return output

 

3. Train

빠른 결과 확인을 위해 epoch=2로 train을 수행했다.

# Instantiate a neural network model 
model = Network()

# train
train(model, 2, train_loader, test_loader, model_name='cifar10_tutorial_epoch2')

The model will be running on cpu device
[1,  1000] loss: 1.777
[1,  2000] loss: 1.487
[1,  3000] loss: 1.334
[1,  4000] loss: 1.235
[1,  5000] loss: 1.158
For epoch 1 the test accuracy over the whole test set is 63 %
[2,  1000] loss: 1.084
[2,  2000] loss: 1.031
[2,  3000] loss: 0.986
[2,  4000] loss: 1.008
[2,  5000] loss: 0.949
For epoch 2 the test accuracy over the whole test set is 67 %

 

4. Test

학습 후 저장된 모델을 불러와서 일부 test 데이터에서 결과를 확인해보자

model.load_state_dict(torch.load('./model/cifar10_tutorial_epoch2.pth'))
testBatch(model, test_loader, batch_size, classes)

 

다음편에서는 동일한 모델을 사용해서 내가 가지고 있는 custom image dataset에 적용해보자.

 

 

참고자료

https://learn.microsoft.com/en-us/windows/ai/windows-ml/tutorials/pytorch-train-model