AlexNet with Pytorch

이 글에서는 pytorch를 이용해서 AlexNet을 구현해보려 합니다. 모델을 학습시키고 평가하는 과정은 제가 저번에 적었던 글에 있는 것을 그대로 다시 사용합니다.

2022.11.14 - [구현] - CNN with Pytorch

CNN with Pytorch

AlexNet 구조 알아보기

Pytorch에서는 수많은 CNN 구조의 딥러닝 모델들을 제공하고 있습니다. AlexNet 모델 역시 pytorch에서 제공하고 있습니다.

import torchvision

alexnet = torchvision.models.alexnet(pretrained=True)
print(alexnet)

pretrained를 True 값으로 하면 ImageNet 데이터셋으로 학습이 미리 진행된 모델을 불러올 수 있습니다. 구조를 보면 아래와 같습니다.

print(alexnet) 한 결과

크게 3가지로 나눠서 볼 수 있을 것 같습니다.

1. features: convolution과 activation fucntion, pooling layer를 연속적으로 거치며 이미지의 특징을 학습하는 구간
2. avgpool: output_size만큼의 출력이 나오도록 average pooling 작업을 거치는 layer
3. classifier: dropout과, linear 연산, activation function을 연속적으로 거치며 이미지를 분류하는 구간

Convolution층들을 거칠 때는 kernel size를 11부터 시작해서 절반씩 줄이다가 3으로 유지하는 것을 볼 수 있고, 활성화 함수로는 ReLU, max pooling 기법을 사용하고 있음을 알 수 있습니다. 사실 이 구조에서 마지막 out_features 1000에 대해 softmax 함수가 적용되는 레이어도 있어야 합니다. 하지만 위의 구조를 보면 softmax 레이어를 딱히 정의하지 않았습니다. 이에 대한 답변은 아래의 사이트에서 찾을 수 있었습니다.

https://discuss.pytorch.org/t/torchvision-models-dont-have-softmax-layer/18071

Torchvision models dont have softmax layer？

Torchvisioin의 모델들은 모두 nn.CrossEntropyLoss()를 사용해 학습된다고 합니다. nn.CrossEntropyLoss()를 보면 log softmax가 포함되어 있습니다. 따라서 모델의 구조를 정의할 때 softmax 층이 없습니다. 그래서 모델을 학습시킬 때 cross entropy loss를 사용해야 합니다.

AlexNet 구조를 변경하고 싶은 경우

torchvision에서 제공하는 모델들은 ImageNet 데이터셋에 대해 학습하는 모델들이기 때문에 최종적으로 1000개의 class에 대한 분류를 진행합니다. 만약, 이 모델을 CIFAR-10 데이터셋에 대해 학습을 진행해야 하는 상황이 온다면 classifier의 마지막 구조를 바꿔야 합니다. 이러한 경우 torchvision에서 가져온 모델의 구조를 바꾸는 방법이 있습니다.

import torchvision
import torch.nn as nn

alexnet = torchvision.models.alexnet(pretrained=True)
in_features = alexnet.classifier[6].in_features
alexnet.classifier[6] = nn.Linear(in_features, 10)

AlexNet에서 사용되는 파라미터의 수

AlexNet에서 사용될 파라미터의 수를 살펴봅시다. Input의 size가 227 * 227로 생각하고 model이 만들어졌기 때문에 이를 생각하시면서 output shape의 변화를 보시면 됩니다.

AlexNet에서 사용되는 파라미터 수

합성곱 층들에서 사용된 파라미터의 수는 (23,296 + 307,392 + 663,936 + 884,992 + 590,080) = 2,469,696개입니다. 얼추 반올림하면 247만 개나 되어 많아 보입니다. 하지만 fully-connected layer들에서 사용되는 파라미터 수는 (37,352,832 + 16,781,312 + 4,097,000) = 58,231,144개로 약 23.58배 많은 것을 알 수 있습니다.

Pytorch로 AlexNet 구조 따라 만들어보기

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


class AlexNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, stride=1, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2)
        )
        self.avg_pool = nn.AdaptiveAvgPool2d(output_size=6)
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, 1000)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avg_pool(x)
        x = x.view(-1, 256 * 6 * 6)
        x = self.classifier(x)
        return x

nn.Sequential을 통해 여러 레이어들을 묶어서 저장할 수 있습니다. 따라서 forward 함수를 작성할 때 이전 글처럼

x = self.pooling(self.relu(self.conv1(x))), x = self.pooling(self.relu(self.conv2(x))),... 형태로 작성할 필요 없이

x = self.features(x) 한 줄로 코드를 깔끔하게 작성할 수 있게 됩니다.