[pytorch] 합성곱 신경망 (Convolution Neural Network) 예제 코드

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.init as init
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

batch_size = 256
learning_rate = 0.0002
num_epoch = 10

mnist_train = dset.MNIST("./", train=True, transform=transforms.ToTensor(), target_transform=None, download=True)
mnist_test = dset.MNIST("./", train=False, transform=transforms.ToTensor(), target_transform=None, download=True)

# DataLoader -> 데이터를 256개씩 묶음, 셔플하고 (test데이터는 안함), 묶을 때 사용할 프로세스 개수는 2개, 묶고 남는건 버림.
train_loader = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=True)
test_loader = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=2, drop_last=True)


# 재사용성을 위해 모델을 클래스로 설계
class CNN(nn.Module):
  # 클래스의 기본값 설정하기
  def __init__(self):
    # CNN클래스의 부모 클래스인 nn.Module 클래스를 반환, CNN클래스의 인스턴스에 사용되는 self를 전달
    super(CNN,self).__init__()
    self.layer = nn.Sequential(
        # nn.Conv2d(in_channel, out_channel (필터 개수), kernel_size)
        # stride, padding = default -> (1,0)
        # data shape -> [batch_size, 1, 28, 28]
        nn.Conv2d(1,16,5), # data shape -> [batch_size, 16, 24, 24]
        nn.ReLU(),
        nn.Conv2d(16,32,5), # data shape -> [batch_size, 32, 20, 20]
        nn.ReLU(),
        # nn.MaxPool2d(kernel_size, stride)
        nn.MaxPool2d(2,2), # data shape -> [batch_size, 32, 10, 10]
        nn.Conv2d(32,64,5), # data shape -> [batch_size, 64, 6, 6]
        nn.ReLU(),
        nn.MaxPool2d(2,2) # data shape -> [batch_size, 64, 3, 3]
    )
    self.fc_layer = nn.Sequential(
        nn.Linear(64*3*3, 100), # 64*3*3개의 노드를 100개로 줄임
        nn.ReLU(),
        nn.Linear(100, 10) # 100개의 노드를 10개로 줄임
    )
  def forward(self,x):
    out = self.layer(x) # data가 합성곱 신경망을 통과
    out = out.view(batch_size, -1) # [batch_size, 64, 3, 3] -> [batch_size, -1]
    out = self.fc_layer(out) # fully connected layer를 통과
    return out


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = CNN().to(device)
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters() , lr=learning_rate)


# Training
loss_arr = []
for i in range(num_epoch):
  # train_loader에서 데이터 소진될때까지 뽑아오고, 인덱스(j)와 [image, label]을 반환
  for j,[image, label] in enumerate(train_loader):
    x = image.to(device)
    y = label.to(device)

    # gradient 누적을 막기 위해 매번 0으로 초기화
    optimizer.zero_grad()
    output = model.forward(x)
    loss = loss_func(output, y)
    # gradient 연산
    loss.backward()
    # 파라미터 업데이트
    optimizer.step()

    if j%1000==0:
      print(loss)
      loss_arr.append(loss.cpu().detach().numpy())


# Test data에 대해 Validation
correct = 0
total = 0
# gradient 연산은 안함
with torch.no_grad():
  for image, label in test_loader:
    x = image.to(device)
    y = label.to(device)

    output = model.forward(x)
    # output 중 최대값을 찾아 그 인덱스를 반환
    _,output_index = torch.max(output,1)

    total +=label.size(0)
    correct +=(output_index==y).sum().float()
  print("Accuracy for Test Data : {}".format(100*correct/total))