[PyTorch] convolutional layer 출력 크기 계산하는 함수 만들기

 CNN 모델을 구축하다보면 conv layer의 출력값 계산을 실수하여 모델이 오류가 발생하는 경우가 종종 있습니다.

 실수를 방지하기 위해 conv layer의 출력값을 계산해주는 함수입니다.

 

import torch.nn as nn
import numpy as np

# define the helper function
def findConv2dOutShape(H_in, W_in, conv, pool=2):
    # get conv arguments
    kernel_size = conv.kernel_size
    stride = conv.stride
    padding = conv.padding
    dilation = conv.dilation

    H_out = np.floor((H_in + 2*padding[0] - dilation[0]*(kernel_size[0]-1)-1) / stride[0] + 1)
    W_out = np.floor((W_in + 2*padding[1] - dilation[1]*(kernel_size[1]-1)-1) / stride[1] + 1)

    if pool:
        H_out /= pool
        W_out /= pool
    
    return int(H_out), int(W_out)

 

 위 계산 식은 PyTorch 공식 문서에 나와있는 식을 그대로 이용했습니다.

 

 함수 인자에 피쳐맵의 높이, 넓이, filter 사이즈, pooling 정보를 입력해주면 출력값을 계산해줍니다.  

 

conv1 = nn.Conv2d(3, 8, kernel_size=3)
h,w = findConv2dOutShape(96,96,conv1)
print(h,w)

 

 

 96x96 이미지를 입력했을 때, 47x47의 피쳐맵이 생성됩니다.

 

 이 함수를 활용해서 간단한 CNN 모델을 구축해보도록 하겠습니다.

 

# implement the CNN model
import torch.nn as nn
import torch.nn.functional as F

# define the Net class
class Net(nn.Module):
    def __init__(self, params):
        super(Net, self).__init__()
        C_in, H_in, W_in = params['input_shape']
        init_f = params['initial_filters']
        num_fc1 = params['num_fc1']
        num_classes = params['num_classes']
        self.dropout_rate = params['dropout_rate']

        self.conv1 = nn.Conv2d(C_in, init_f, kernel_size=3)
        h,w = findConv2dOutShape(H_in, W_in, self.conv1)
        self.conv2 = nn.Conv2d(init_f, 2*init_f, kernel_size=3)
        h,w = findConv2dOutShape(h, w, self.conv2)
        self.conv3 = nn.Conv2d(2*init_f, 4*init_f, kernel_size=3)
        h,w = findConv2dOutShape(h, w, self.conv3)
        self.conv4 = nn.Conv2d(4*init_f, 8*init_f, kernel_size=3)
        h,w = findConv2dOutShape(h, w, self.conv4)

        # computer the flatten size
        self.num_flatten = h*w*8*init_f
        self.fc1 = nn.Linear(self.num_flatten, num_fc1)
        self.fc2 = nn.Linear(num_fc1, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv3(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv4(x))
        x = F.max_pool2d(x, 2, 2)

        x = x.view(-1, self.num_flatten)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, self.dropout_rate, training = self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

 

 최종 conv layer에서 생성된 출력값은 fc layer에 입력할 때 크기를 변경시켜줘야 합니다.

 위에서 정의한 함수를 사용하면, 입력 이미지가 conv layer를 거쳐서 fc layer에 입력될 때, 크기가 자동으로 변경됩니다.