[머신러닝] CNN 모델 구현 with Pytorch (CIFAR-10 dataset)

2022.07.21 - [Studying/Machine Learning] - [머신러닝] Convolution Layer - Padding, Stride, Dilation

[머신러닝] Convolution Layer - Padding, Stride, Dilation

이전 포스트까지 CNN에 대한 것들을 대부분 다룬 것 같다.

 

이제 파이썬의 Pytorch로 사용해보겠다.

 

---

데이터는 torchvision.datasets에 있는 CIFAR10을 사용하겠다.

해당 데이터 셋은 batch_size * 3 * 32 * 32의 사이즈를 가지고 있다.

 

이 dataset은 아래 sample과 같이 32*32픽셀의 60000개 컬러 이미지가 있고 총 10개의 클래스로 라벨링 되어있다.

우선 import를 해온다.

import torch
import torch.nn as nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms

 

CNN 모델을 구현한다.

 

우리는 batch, 3, 32, 32 사이즈의 인풋을 받아서

batch, 10의 아웃풋을 낼것이다.

 

pytorch의 Conv2d layer를 사용하려 하는데 파라미터는 다음과 같다.

in_channel, out_channel, kernel, stride, padding, dilation

위 파라미터의 개념은 이전 포스트에서 다뤘다.

 

더 자세한 옵션은 아래 링크를 참고하면 된다.

https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html

 

각 레이어 통과 후 사이즈는 코드 내부에 주석으로 써놓겠다.

class CNN_prac(nn.Module):
    def __init__(self):
        super(CNN_prac, self).__init__()
       
        #input -> N, 3, 32, 32
        self.conv1 = nn.Conv2d(3, 512, 3, 1, 1)
        #output -> N, 512, 32, 32(batch, output, h, w)
        
        self.bn = nn.BatchNorm2d(512)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.1)

        self.conv2 = nn.Conv2d(512, 256, 3, 1, 1) 
        #output -> N, 256, 32, 32(batch, output, h, w)        
        self.bn2 = nn.BatchNorm2d(256)

        self.conv3 = nn.Conv2d(256, 256, 3, 2, 1)
        #output -> N, 256, 16, 16(batch, output, h, w)
        self.bn3 = nn.BatchNorm2d(256)
        

        self.conv4 = nn.Conv2d(256, 256, 3,4,1)
        #output -> N, 256, 4, 4(batch, output, h, w)
        self.bn4 = nn.BatchNorm2d(256)

        self.linear = nn.Linear(256*4*4, 10)
        #output -> N, 10

    def forward(self, x):
        out=self.conv1(x)
        out=self.bn(out)
        out = self.relu(out)
        out = self.dropout(out)

        out=self.conv2(out)
        out=self.bn2(out)
        out = self.relu(out)
        out = self.dropout(out)

        out=self.conv3(out)
        out=self.bn3(out)
        out = self.relu(out)
        out = self.dropout(out)

        out=self.conv4(out)
        out=self.bn4(out)
        out = self.relu(out)
        out = self.dropout(out)

        # (N, 256, 4, 4) -> ( N, 256x4x4 )
        out = out.reshape(-1, 256*4*4)
        
        out = self.linear(out)
        return out

 

이제 모델을 제대로 구현했는지 확인하기 위해 sample을 만들어 테스트를 해보겠다.

 

cnn_prac=CNN_prac()
sample_image = torch.zeros(64, 3, 32, 32)
output=cnn_prac(sample_image)
print(output.size())
#torch.Size([64, 10])

print(cnn_prac)
# CNN_prac(
#   (conv1): Conv2d(3, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#   (relu): ReLU()
#   (dropout): Dropout(p=0.1, inplace=False)
#   (conv2): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#   (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#   (conv3): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
#   (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#   (conv4): Conv2d(256, 256, kernel_size=(3, 3), stride=(4, 4), padding=(1, 1))
#   (bn4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#   (linear): Linear(in_features=4096, out_features=10, bias=True)
# )

output이 64,10으로 잘 나왔고

모델도 원하는 대로 구현이 잘 된 것을 볼 수 있다.

 

---

이제 데이터를 불러오겠다.

이미지는 데이터가 한정적으로 있기 때문에 약간의 트릭을 사용하겠다.

transform을 이용해서 이미지를 회전시키고 뒤집고 해서 몇 가지의 이미지를 더 만들어 보겠다.

이를 Data Augmentation이라고 한다.

 

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

더 많은 옵션은 위 링크를 참고하면 된다.

 

train_dataset = datasets.CIFAR10('./',
                 train=True,
                 transform=transforms.Compose(
                     [
                        transforms.RandomHorizontalFlip(0.5),
                        transforms.RandomRotation(30),
                        transforms.Resize((32,32)),
                        transforms.ToTensor()
                     ]
                 ),
                 download=True)

test_dataset = datasets.CIFAR10('./',
                 train=True,
                 transform=transforms.Compose(
                     [
                        transforms.Resize((32,32)),
                        transforms.ToTensor()
                     ]
                 ),
                 download=True)

 

데이터 다운이 완료되었다.

 

---

Training을 시키기 전에 numpy 패키지를 가져오고

import numpy as np

 

 

이미지는 연산이 오래 걸리기 때문에 모델을 GPU로 올려서 계산하도록 cuda()를 사용해주겠다.

 

loss는 pytorch에서 제공하는 CrossEntropyLoss를 사용하고,

Adam optimizer를 사용하겠다.

 

max_patience = 5 
patience = 0
best_loss = 1000

model = CNN_prac()
model = model.cuda() # cpu -> gpu
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

batch_size = 128
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

 

이제 학습을 시켜보겠다.

코드에 주석으로 설명을 달아놓았다.

num_epochs = 10

for epoch in range(num_epochs):
################# Train #################
  model.train()
  train_loss = []
  for batch in train_dataloader:
    image, label = batch
    image = image.cuda()
    label = label.cuda()#data를 cpu -> gpu로 올리기

    output = model(image) 

    loss = criterion(output, label) 

    optimizer.zero_grad()#이전 gradient값 없애기
    loss.backward()

    optimizer.step()#update

    train_loss.append(loss.item())

  print('Epoch {} - mean train loss {}'.format(epoch, np.mean(train_loss)))


################# Test #################
  model.eval() # 사용 이유 뒤에 설명
  with torch.no_grad():
  # test기 때문에 gradient를 계산하지 않아 빨라지고 메모리 사용량이 감소
    test_loss = []
    for batch in test_dataloader:
      image, label = batch
      image = image.cuda()
      label = label.cuda()

      output = model(image) 

      loss = criterion(output, label) 

      test_loss.append(loss.item())
  print('Epoch {} - mean test loss {}'.format(epoch, np.mean(test_loss)))

  if best_loss > np.mean(test_loss):
      best_loss = np.mean(test_loss)
      torch.save(model.state_dict(), "./model.pth")
      #현재 모델을 저장한다.
      patience = 0
      print("model saved - after epoch {}".format(epoch))
  else:
    if patience == max_patience:
        patience +=1
        break

 

model.eval()을 사용하는 이유는

학습시킬 때와 테스트할 때 다르게 동작해야 하는 부분이 있기 때문이다.

 

주로 dropout layer나 batch norm layer인데 이 layer들은 테스트 때 사용하지 않는다.

각 layer를 간단하게 설명하면

 

dropout은 training시 overfitting을 해소하기 위해 일부 뉴런들을 drop 시키는 것이다.

테스트에서는 당연히 사용할 필요가 없다.

 

batch normalization은 batch 안에서 정규화를 해서 local optimum에 빠지지 않도록 도와주는 것인데

local, global optimum 부분도 다음에 정리해보겠다.

 

---

이렇게 CNN 모델을 구현하고 학습시키고 학습된 모델을 저장해봤다.

 

다음 포스트에는 local/global optimum을 정리하거나 기 학습된 CNN모델을 사용해보겠다.