[논문 구현] ShuffleNet v1 직접 구현해보기

이번에는 ShuffleNet v1을 직접 구현해보았다. 확실히 MobileNet v1과 비교했을 때 비슷한 학습속도에서 높은 성능을 보여줌을 확인할 수 있었다.

 

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. GConv

 

ShuffleNet unit에 Group Convolution이 적용되며 unit에는 2개의 GConv가 존재한다. 첫번째 GConv는 채널 셔플이 적용되고 두번째는 적용되지 않기에 그것을 구분하는 인자를 두어 ReLU 적용 여부와 Channel Shuffle여부를 체크하도록 했다. 또한 Stage2에서는 첫번째 GConv에서 Channel Shuffle을 적용하지 않으며 input channel과 output channel의 차이가 2배가 아닌 큰 차이가 나타나므로 뒤에 코드에서 다시 언급하겠지만 skip connection 적용시 1x1 Conv로 channel dimension을 조정한 후 concat 해주는 형태로 dimension을 조정했다. 정답은 아니므로 참고만 하면 좋을 것 같다.

 

또한 Channel shuffle은 논문에서 언급한 방법대로 reshape / transpose / flatten 순으로 적용하였다.

논문에서 언급한 Channel Shuffle 방법

그림 (c) 처럼 channel shuffle이 이루어지며 모든 channel을 연관시키는 방법으로 논문에서 언급하고 있다.

class GConv(nn.Module):
    def __init__(self, groups, in_channels, out_channels, channel_shuffle = True, s2 = False):
        super(GConv, self).__init__()
        self.s2 = s2
        self.channel_shuffle = channel_shuffle
        self.groups = groups
        self.out_channels = out_channels
        self.gconv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=self.out_channels,
            groups=self.groups,
            kernel_size=1,
            stride=1,
            padding=0
        )
        if s2:
            self.gconv = nn.Conv2d(
                in_channels=in_channels,
                out_channels=self.out_channels,
                kernel_size=1,
                stride=1,
                padding=0
            )
            self.channel_shuffle = False
        self.batch = nn.BatchNorm2d(num_features=self.out_channels)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.gconv(x)
        x = self.batch(x)
        if self.channel_shuffle:
            x_shape = x.shape
            x = self.relu(x)
            x = x.reshape(x_shape[0], self.groups, self.out_channels//self.groups, x_shape[-2], -1)
            x = torch.transpose(x, 1, 2)
            x = torch.flatten(x, start_dim=1, end_dim=2)
        if self.s2:
            x = self.relu(x)
        return x

 

3. DWConv

ShuffleNet Unit에는 DWConv도 적용되므로 클래스를 선언해주었다.

 

class DWConv(nn.Module):
    def __init__(self, in_channels, stride, multiplier = 1):
        super(DWConv, self).__init__()
        out_channels = in_channels * multiplier
        self.net = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            groups=in_channels,
            kernel_size=3,
            padding=1,
            stride=stride
        )
        self.batch = nn.BatchNorm2d(num_features=out_channels)

    def forward(self, x):
        x = self.net(x)
        x = self.batch(x)
        return x

 

4. ShuffleNet Unit

 

앞의 GConv에서 언급한대로 Stage2에서는 stride2인 첫 unit에서 input dimension과 output dimension 차이가 다른 Stage간 차이와 다르므로 identity Conv를 활용해 dimension을 맞춰준 후 concat이 되도록했다. 나머지 stride2 인 부분은 그대로 avgpooling만 통과한 후 net통과한 값과 concat하는 연산을 했고, stride1인 부분은 pooling 없이 add 형태로 연산을 하도록 했다.

 

class ShuffleNetUnit(nn.Module):
    def __init__(self, stride, in_channels, out_channels, groups, s2=False):
        super(ShuffleNetUnit, self).__init__()
        self.s2 = s2
        self.stride = stride
        if stride == 1:
            d = 1
        else:
            d = 2
            self.avgpool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)
        out_channels = out_channels // d
        self.gconv1 = GConv(groups=groups, in_channels=in_channels, out_channels=out_channels, s2=s2)
        self.DWConv = DWConv(in_channels=out_channels, stride = stride)
        self.gconv2 = GConv(groups=groups, in_channels=out_channels, out_channels=out_channels, channel_shuffle=False)
        if s2:
            self.identity = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0)
        self.relu = nn.ReLU()

    def forward(self, x):
        x2 = self.gconv1(x)
        x2 = self.DWConv(x2)
        x2 = self.gconv2(x2)
        if self.stride == 2:
            x = self.avgpool(x)
            if self.s2:
                x = self.identity(x)
            x = torch.concat((x, x2), dim=1)
        else:
            x = torch.add(x, x2)
        x = self.relu(x)
        return x

 

5. ShuffleNet v1

 

논문에서 제시한 전체 아키텍쳐를 groups 인자에 따라 다르게 적용되도록 만들었다.

 

class ShuffleNetV1(nn.Module):
    def __init__(self, in_channels, groups, num_classes):
        super(ShuffleNetV1, self).__init__()
        self.channels_dict = {
            1 : 144,
            2 : 200,
            3 : 240,
            4 : 272,
            8 : 384
        }
        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=24,
            kernel_size=3,
            stride=2,
            padding=1
        )
        self.maxpool = nn.MaxPool2d(
            kernel_size=3,
            stride=2,
            padding=1
        )
        channel_var = self.channels_dict[groups]
        stage2_list = [
            ShuffleNetUnit(
                stride=1,
                in_channels=channel_var,
                out_channels=channel_var,
                groups=groups
            ) for _ in range(3)
        ]
        s2 = [
                 ShuffleNetUnit(
                     stride=2,
                     in_channels=24,
                     out_channels=channel_var,
                     groups=groups,
                     s2=True

            )] + stage2_list
        self.stage2 = nn.Sequential(*s2)
        stage3_list = [
            ShuffleNetUnit(
                stride=1,
                in_channels=channel_var * 2,
                out_channels=channel_var * 2,
                groups=groups
            ) for _ in range(7)
        ]
        s3 = [
            ShuffleNetUnit(
                stride=2,
                in_channels=channel_var,
                out_channels=channel_var * 2,
                groups=groups
            )
        ] + stage3_list
        self.stage3 = nn.Sequential(*s3)
        stage4_list = [
            ShuffleNetUnit(
                stride=1,
                in_channels=channel_var * 4,
                out_channels=channel_var * 4,
                groups=groups
            ) for _ in range(3)
        ]
        s4 = [
            ShuffleNetUnit(
                stride=2,
                in_channels=channel_var * 2,
                out_channels=channel_var * 4,
                groups=groups
            )
        ] + stage4_list
        self.stage4 = nn.Sequential(*s4)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(in_features=channel_var*4, out_features=num_classes, bias=True)
        self.soft = nn.Softmax(dim = 1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.maxpool(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = self.avgpool(x)
        x = x.squeeze()
        x = self.fc(x)
        x = self.soft(x)
        return x

 

CIFAR-10로 테스트한 결과 MobileNet v1과 비교했을 때 같은 Epoch 20에서 비슷한 학습속도로 더 큰 성능을 나타냈다.

MobileNet v1 : 0.553

ShuffleNet v1 : 0.638

 

참고로 ResNet101은 0.691이지만 학습속도가 2배이상이다.

 

전체 코드는 다음 링크를 참고하면된다.

https://github.com/kkt4828/reviewpaper/blob/68b179ff6550744ea262306908be50eb386dd3cc/ShuffleNet/ShuffleNetV1.py

GitHub - kkt4828/reviewpaper: 논문구현

 

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

ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices