[Pytorch] CNN 모델 돌려보기 예제

0. 글 소개

• 딥러닝 학습을 위한 기초 예제코드
• 출처: pytorch 튜토리얼

1. 모델 클래스 정의

import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__() # 부모 클래스 상속 받아야함
        
        # 컨볼루션 커널 정의
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 아핀(affine) 연산: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5*5은 이미지 차원에 해당
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
    	"""
        
        backward함수는 자동으로 정의되어 만들지 않아도 됨
        """
    
        # (2, 2) 크기 윈도우에 대해 맥스 풀링(max pooling)
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # 크기가 제곱수라면, 하나의 숫자만을 특정(specify)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1) # 배치 차원을 제외한 모든 차원을 하나로 평탄화(flatten)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()
print(net)

 

• in_channels: input 데이터의 채널 수
• out_channels: output 데이터의 채널 수
• kernel_size:  사이즈 

torch.nn.Conv2d(
    in_channels, 
    out_channels, 
    kernel_size, 
    stride=1, 
    padding=0, 
    dilation=1, 
    groups=1, 
    bias=True, 
    padding_mode='zeros'
)

• in_features: 입력되는 x의 차원
• out_features: 출력되는 y의 차원

torch.nn.Linear(
    in_features,
    out_features,
    bias=True,
    device=None,
    dtype=None
)

 

2. optimizer, loss function 정의

• 까먹지말고 꼭 해줄 것!!!

cnn = CNN()
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(cnn.parameters(), lr=0.01)

3. 학습

• 모든 단계에서 빠지는 것이 있지 않도록....!!

cnn.train()  # 학습 부분
for epoch in range(10):
  for index, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()  # 기울기 초기화, 변화도 버퍼를 0으로
    output = cnn(data) # 예측
    loss = criterion(output, target) # loss 계산
    loss.backward()  # 역전파 계산
    optimizer.step()  # 옵티마이저 업데이트 진행