[논문 구현] MobileNet V2 직접 구현해보기!!!

 

논문에서 제시한 내용만 보고 구현한 코드

 

약간의 Naive한 점이 있지만 확실히 ResNet 등의 CNN 모델보다 학습속도가 빠르다!

 

그럼 바로 시작

 

1. Setup

 

기존과 동일하다.

 

import torch.nn as nn
import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.optim as optim
import time
import numpy as np

import random
import torch.backends.cudnn as cudnn

seed = 2022
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
cudnn.benchmark = False
cudnn.deterministic = True
random.seed(seed)

2. BottleNeck Residual Block

 

stride 1에 대해서 skip connection을 추가해줬다.

 

논문에서 multiplier는 6으로 고정시켰으나 상황에따라 조정할 수 있게 parameter로 넣어두었다.

 

layer마다 BN은 추가하여 실험했기에 내 코드에도 추가해두었다.

 

class Bottleneck(nn.Module):
    def __init__(self, in_channels, out_channels, multiplier = 6, stride = 2):
        super(Bottleneck, self).__init__()
        self.stride = stride
        mid_ch = in_channels * multiplier
        self.net = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=mid_ch,
                kernel_size=1,
                stride=1
            ),
            nn.BatchNorm2d(mid_ch),
            nn.ReLU6(),
            nn.Conv2d(
                in_channels=mid_ch,
                out_channels=mid_ch,
                stride=stride,
                kernel_size=3,
                padding=1,
                groups=mid_ch,
            ),
            nn.BatchNorm2d(mid_ch),
            nn.ReLU6(),
            nn.Conv2d(
                in_channels=mid_ch,
                out_channels=out_channels,
                stride=1,
                kernel_size=1
            ),
            nn.BatchNorm2d(out_channels)
        )
        self.identity = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=1,
            stride=1
        )
        self.relu = nn.ReLU6()

    def forward(self, x):
        x2 = self.net(x)
        if self.stride == 1:
            x = self.identity(x)
            x2 = torch.add(x2, x)
        x2 = self.relu(x2)
        return x2

 

 

 

3. MobileNet v2

 

Block을 반복해서 넣어야하는 부분이 귀찮아 함수로 만들었다.

 

def make_layer(n, in_channels, out_channels, stride, multiplier):
    b_l = []
    for i in range(n):
        if i == 0:
            b_l.append(Bottleneck(
                in_channels=in_channels,
                out_channels=out_channels,
                stride=stride,
                multiplier=1
            ))
        else:
            b_l.append(
                Bottleneck(
                    in_channels=out_channels,
                    out_channels=out_channels,
                    stride=1,
                    multiplier=multiplier
                )
            )
    return b_l

 

 

전체 Architecture에 맞게 구성하면 다음과 같이 코드를 짤 수 있다.

 

class MobileNetV2(nn.Module):
    def __init__(self, in_channels, num_classes, multiplier=6):
        super(MobileNetV2, self).__init__()
        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=32,
            kernel_size=3,
            padding=1,
            stride=2
        )
        self.relu1 = nn.ReLU6()
        self.batch1 = nn.BatchNorm2d(32)
        self.bottleneck1 = Bottleneck(
            in_channels=32,
            out_channels=16,
            multiplier=1,
            stride=1
        )
        self.bottleneck2 = nn.Sequential(*make_layer(2, 16, 24, 2, multiplier))
        self.bottleneck3 = nn.Sequential(*make_layer(3, 24, 32, 2, multiplier))
        self.bottleneck4 = nn.Sequential(*make_layer(4, 32, 64, 2, multiplier))
        self.bottleneck5 = nn.Sequential(*make_layer(3, 64, 96, 1, multiplier))
        self.bottleneck6 = nn.Sequential(*make_layer(3, 96, 160, 2, multiplier))
        self.bottleneck7 = nn.Sequential(*make_layer(1, 160, 320, 1, multiplier))

        self.conv2 = nn.Conv2d(
            in_channels=320,
            out_channels=1280,
            kernel_size=1,
            stride=1
        )
        self.batch2 = nn.BatchNorm2d(1280)
        self.relu2 = nn.ReLU6()
        self.avgpool = nn.AdaptiveAvgPool2d(1)

        self.conv3 = nn.Conv2d(
            in_channels=1280,
            out_channels=num_classes,
            kernel_size=1,
            stride=1
        )
        self.soft = nn.Softmax(dim = 1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.batch1(x)
        x = self.relu1(x)
        x = self.bottleneck1(x)
        x = self.bottleneck2(x)
        x = self.bottleneck3(x)
        x = self.bottleneck4(x)
        x = self.bottleneck5(x)
        x = self.bottleneck6(x)
        x = self.bottleneck7(x)

        x = self.conv2(x)
        x = self.batch2(x)
        x = self.relu2(x)
        x = self.avgpool(x)
        x = self.conv3(x)
        x = x.squeeze()
        x = self.soft(x)
        return x

 

첫 bottleneck은 multiplier가 1이라 직접 입력해놓았고, nn.Linear가 아닌 1x1 Conv와 squeeze를 활용해 마지막 classification을 위한 tensor를 구성한 후 softmax에 통과시켜 return 시켰다.

 

모든 Activation Function은 ReLU6로 사용했다.

 

논문리뷰는 다음 링크를 참고하면 좋을 것이다.

https://ai.dreamkkt.com/60

[논문 리뷰] MobileNet V2 간단 리뷰!!

참고 논문 - https://arxiv.org/abs/1801.04381v4

MobileNetV2: Inverted Residuals and Linear Bottlenecks