[논문 리뷰 및 구현] MobileNet V3 간단하게 리뷰하고 구현해보기!!

 

핵심아이디어는 SENet에서 나왔던 channel-wise response를 계산하여 이를 적용하는 것과 ReLU6에 H-Swish(HardSwish) 이다.

 

1. H-Swish

 

$$ Swish(x) = x \sigma (x) $$

로 표현되며 기존 ReLU에서 음수값에 대한 정보손실 문제를 해결하기 위한 함수이다.

 

하지만 Sigmoid 자체가 연산량이 크기 때문에 V3에서는 이를 줄이기 위해 Approximation을 한 함수를 제안한다.

 

$$ H-Swish (x) = x \frac{ReLU6(x+3)}{6} $$

 

이는 Swish를 잘 approximation 하고, 연산량도 줄어들어 mobile network에 필요한 연산량을 줄이는 방향과 부합한 방법이다.

 

그림에서 나온 것처럼 거의 같다.

 

 

 

2. V3 Unit

 

 

위의 그림처럼 ReLU6 대신 H-Swish로 대체하는데 전부다 하지는 않고 어느정도의 depth 이후에 적용하는 식으로 설계했다.

 

feature 추출 시 바로 1x1 projection 통과 전 Squeeze-extraction(SE)을 적용하는 방법으로 구성되었다.

 

 

 

구현 코드는 다음과 같다.

 

class SqueezeExcitation(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(SqueezeExcitation, self).__init__()
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels // 4,
            kernel_size=1,
            stride=1
        )
        self.batch1 = nn.BatchNorm2d(out_channels // 4)
        self.relu = nn.ReLU()
        self.fc2 = nn.Conv2d(
            in_channels=out_channels // 4,
            out_channels=out_channels,
            kernel_size=1,
            stride=1
        )
        self.batch2 = nn.BatchNorm2d(out_channels)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.avgpool(x)
        x = self.fc1(x)
        x = self.batch1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.batch2(x)
        x = self.sigmoid(x)
        return x

SENet에 있던 squeeze - extraction expansion layer를 out channel의 1/4로 수정했다.

 

class bottleneck(nn.Module):
    def __init__(self, in_channels, out_channels, mid_ch, stride = 2, se=False, size=3, mode = 'RE'):
        super(bottleneck, self).__init__()
        padding = 1
        if size == 3:
            padding = 1
        elif size == 5:
            padding = 2
        self.se = se
        self.stride = stride
        self.conv1 = nn.Conv2d(
            in_channels = in_channels,
            out_channels = mid_ch,
            kernel_size=1,
            stride=1
        )
        self.batch1 = nn.BatchNorm2d(mid_ch)
        if mode == 'RE':
            self.act1, self.act2, self.act3 = nn.ReLU6(), nn.ReLU6(), nn.ReLU6()
        elif mode == 'HS':
            self.act1, self.act2, self.act3 = nn.Hardswish(), nn.Hardswish(), nn.Hardswish()

        self.DWConv = nn.Conv2d(
            in_channels=mid_ch,
            out_channels=mid_ch,
            groups=mid_ch,
            kernel_size=size,
            stride=stride,
            padding=padding
        )
        self.batch2 = nn.BatchNorm2d(mid_ch)

        self.seblock = SqueezeExcitation(
            in_channels=mid_ch,
            out_channels=mid_ch
        )

        self.conv2 = nn.Conv2d(
            in_channels=mid_ch,
            out_channels=out_channels,
            kernel_size=1,
            stride=1
        )
        self.batch3 = nn.BatchNorm2d(out_channels)
        if self.stride == 1:
            self.identity = nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=1,
                stride=self.stride
            )

    def forward(self, x):
        x2 = self.conv1(x)
        x2 = self.batch1(x2)
        x2 = self.act1(x2)
        x2 = self.DWConv(x2)
        x2 = self.batch2(x2)
        x2 = self.act2(x2)
        if self.se:
            x3 = self.seblock(x2)
            x2 = torch.mul(x2, x3)
        x2 = self.conv2(x2)
        x2 = self.batch3(x2)
        if self.stride == 1:
            x = self.identity(x)
            x2 = torch.add(x, x2)

        x2 = self.act3(x2)
        return x2

 

RE, HS 모드에 따라 activation이 다르게 적용되도록 구성했고, SE layer 적용 여부를 parameter로 추가해둬 전체 네트워크 구성시 편하게 코딩할 수 있게했다. size가 3x3과 5x5가 섞여서 적용되기에 size parameter도 추가해뒀다.

 

 

 

3. Architecture

 

Large모델에서 마지막 conv2d 1x1 은 BatchNorm을 적용하지 않는 것이 특이한 점이며, 중간부터 H-Swish를 적용한 것도 특징이라 할 수 있다. 나머지 Layer에서는 Batch Norm을 적용했다.

 

small모델은 depth와 width를 줄인 정도이며 비슷한 구조를 갖는다.

 

구현은 Large 모델로 진행했다.

 

class MobileNetV3_Large(nn.Module):
    def __init__(self, num_classes, in_channels):
        super(MobileNetV3_Large, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=16,
                kernel_size=3,
                stride=2,
                padding=1
            ),
            nn.BatchNorm2d(16),
            nn.Hardswish()
        )
        self.bneck1 = bottleneck(in_channels=16, out_channels=16, mid_ch=16, stride=1, mode='RE')
        self.bneck2 = nn.Sequential(
            bottleneck(in_channels=16, out_channels=24, mid_ch=64, stride=2, mode='RE'),
            bottleneck(in_channels=24, out_channels=24, mid_ch=72, stride=1, mode='RE')
        )
        self.bneck3 = nn.Sequential(
            bottleneck(in_channels=24, out_channels=40, mid_ch=72, stride=2, mode='RE', size=5, se=True),
            bottleneck(in_channels=40, out_channels=40, mid_ch=120, stride=1, mode='RE', size=5, se=True),
            bottleneck(in_channels=40, out_channels=40, mid_ch=120, stride=1, mode='RE', size=5, se=True)
        )
        self.bneck4 = nn.Sequential(
            bottleneck(in_channels=40, out_channels=80, mid_ch=240, stride=2, mode='HS'),
            bottleneck(in_channels=80, out_channels=80, mid_ch=200, stride=1, mode='HS'),
            bottleneck(in_channels=80, out_channels=80, mid_ch=184, stride=1, mode='HS'),
            bottleneck(in_channels=80, out_channels=80, mid_ch=184, stride=1, mode='HS')
        )
        self.bneck5 = nn.Sequential(
            bottleneck(in_channels=80, out_channels=112, mid_ch=480, stride=1, mode='HS', se=True),
            bottleneck(in_channels=112, out_channels=112, mid_ch=672, stride=1, mode='HS', se=True)
        )
        self.bneck6 = nn.Sequential(
            bottleneck(in_channels=112, out_channels=160, mid_ch=672, stride=2, mode='HS', se=True, size=5),
            bottleneck(in_channels=160, out_channels=160, mid_ch=960, stride=1, mode='HS', se=True, size=5),
            bottleneck(in_channels=160, out_channels=160, mid_ch=960, stride=1, mode='HS', se=True, size=5)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(
                in_channels=160,
                out_channels=960,
                kernel_size=1,
                stride=1
            ),
            nn.BatchNorm2d(960),
            nn.Hardswish()
        )
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=960, out_channels=1280, kernel_size=1, stride=1),
            nn.Hardswish()
        )
        self.fc = 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.bneck1(x)
        x = self.bneck2(x)
        x = self.bneck3(x)
        x = self.bneck4(x)
        x = self.bneck5(x)
        x = self.bneck6(x)
        x = self.conv2(x)
        x = self.avgpool(x)
        x = self.conv3(x)
        x = self.fc(x)
        x = x.squeeze()
        x = self.soft(x)

        return x

 

마지막 fc layer와 conv2에 Batch Norm을 적용하지 않았다.

 

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