CNN 모델의 classifier단에서 FC layer의 입력 노드 개수

﹡

 

이미지 다중분류 문제에서 CNN을 사용하면 convolution layer 와 ReLU, Batchnorm2d, MaxPool2d를 반복하다가

마지막에는 fully connected layer를 통과해서 softmax를 얻어 분류를 해주어야 한다.

 

예를 들어 CIFAR10 데이터셋으로 10가지를 분류하는 네트워크를 만든다고 하자.

통상적인 네트웍에서는 CNN 레이어를 통과해서 다양한 특징맵을 뽑아내고 풀링을 해서 사이즈를 반으로 줄인다.

(Maxpooling을 한다고 해서 채널수가 줄지는 않는다. 채널별로 각각 풀링해주는 거니까)

 

충분히 특징맵을 얻었다면 nn.Linear( )를 활용해서

최종적인 출력 노드를 10개로 줄여줘야 할 것이다. 

이때 nn.Linear(?, 10)의 ? 부분에 어떤것을 써야할 지 의문이 들었다. 

conv layer를 통과한 직후에 특징맵의 채널수와 행, 열 크기를 한눈에 봐서는 알 수 없기 때문이다. 

이것을 확인하기 위한 방법은 nn.Linear() 부분을 주석처리 해놓고 그 전까지 통과한 결과의 shape을 확인해보는 것이다. 

 

import torch
from torch import nn
class CNN(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(1,8,6,stride=2)
        self.conv2 = nn.Conv2d(8,16,3,padding=1)
        self.Maxpool2 = nn.MaxPool2d(2)
        # self.fc = nn.Linear(?,10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.Maxpool2(x)
        return x

x = torch.randn(32,1,28,28)
model = CNN()
print(model(x).shape)

torch.Size([32, 16, 6, 6])

위에서 확인할 수 있듯이 그 결과는 torch.size([32, 16, 6, 6]) 이었다.

즉, fc layer를 통과하기 직전까지 16개의 채널을 가진 6x6 크기의 feature map 32개가 출력되었다는 의미이다. 

데이터 배치 사이즈인 32를 제외하고 16*6*6 은 flatten 시켜주어 바로 nn.Linear()의 입력 노드 개수로서 입력해줘야 한다. 

따라서 ? 에 들어갈 숫자는 16*6*6 이 됨을 알 수 있다. 

 

코드를 수정시켜주어 클래스를 완성시켜주면 다음과 같다. 

fc 레이어에 nn.Linear() 를 추가해주고 x를 flatten 시켜주었다.

import torch
from torch import nn
class CNN(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(1,8,6,stride=2)
        self.conv2 = nn.Conv2d(8,16,3,padding=1)
        self.Maxpool2 = nn.MaxPool2d(2)
        self.fc = nn.Linear(16*6*6,10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.Maxpool2(x)
        x = torch.flatten(x, start_dim=1)
        x = self.fc(x)
        return x

 

* 참고하기

 

https://pytorch.org/docs/stable/generated/torch.flatten.html