[Pytorch] ResNet을 만들어보자!

Pytorch를 이용해서 분류문제에서 사용하는 여러가지 backbone 논문을 구현하는 시간을 가지려고 한다. 네트워크에 대한 논문 리뷰는 따로 하지 않고 구현을 하면서 간단한 설명만 할 예정이다.

 

$\triangledown$ Resnet paper

https://arxiv.org/abs/1512.03385

Deep Residual Learning for Image Recognition

---

구현할 네트워크에 대한 Architecture 정보가 논문에 나와있어 구현하기가 매우 용이하였다. 

 

이제 차례 차례 구현해보자!

Resnet의 첫번째 논문에서는 원래 Convolution output과 ReLu이전의 output간의 skip connection(a)이 이루어졌다. 하지만 두번째 논문에서 이러한 구조가 제대로된 skip connection이 아니라 하고, 여러 실험을 통해 (e)에서 제안하는 순서의 residual connection이 가장 이상 구조라고 한다. 따라서 이번 구현에서는 full pre-activation을 사용할 것이다.

 

'''Base Convolution Block: Bn + Relu + Conv'''
class BnConv(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, stride=stride,
                              padding=padding)
        self.bn = nn.BatchNorm2d(in_ch)

    def forward(self, x):
        return self.conv(F.relu(self.bn(x)))

간편하게 이를 사용하기 위해 다음과 같이 BatchNorm + Relu + Conv Block을 따로 선언하였다. 

 

Stem

나는 Network를 구현할 때, Stem과 Body, Header로 나누는 편이다. 개인적으로 이렇게 분류해놓는 것이 보기나 분석하기 좋다고 생각한다. 먼저 Stem은 Residual Block을 쓰기 이전에 초기 layer를 말한다. 

        self.conv1 = nn.Conv2d(3, first_ch, kernel_size=7, stride=2, padding=3)
        self.pool1 = nn.MaxPool2d(3, 2, padding=1)

논문에서는 7x7 stride 2 Convolution layer와 Maxpool으로 stem을 구성해 놓았다. 초기 이미지의 resolution이 크기 때문에 stride나 pool을 통해서 이를 줄여주는 과정을 거치게 된다. max pool과 stride 2 중 어느 것이 좋은가에 대한 논란도 있고 이에 대한 논문도 있는데 최근에는 stride 2를 주로 사용하는 추세이다. 

 

개인적으로 pytorch를 사용하면서 keras와 다르게 input channel이나 padding 크기를 내가 지정해줘야하는데, 이 때문에 output size를 항상 생각해줄 필요가 있다.

$$OWH = \frac{HW - CHW + 2P}{S} + 1 $$

 

Body

이 부분은 Network의 크기에 따라서 block이 달라진다. 보통 resnet 18, 34는 basic block, 그 이후는 Bottleneck block을 사용한다.

Basic Residual Block

class ResBlock(Module):
    expansion = 1  # out_ch expansion from in_ch

    def __init__(self, in_ch, out_ch, stride=1):
        super().__init__()
        expans_ch = out_ch * self.expansion
        self.conv1 = BnConv(in_ch, out_ch, stride=stride)
        self.conv2 = BnConv(out_ch, out_ch)

        if stride != 1 or in_ch != expans_ch:
            self.residual = BnConv(in_ch, expans_ch, kernel_size=1,
                                   stride=stride, padding=0)
        else:
            self.residual = nn.Sequential()

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        return out + self.residual(x)

expansion이란 변수는 나중에 나올 BottleNeck Block과의 차이를 두기 위함이다. 기본적으로 Basic Block에서는 1이기 때문에 무시하여도 된다. 

Basic Block에서는 Convolution layer 2개를 거친다음에의 output과 input을 더해줌으로써 skip connection이 이루어진다. 이 때 주의해야할 점은 stride에 의해서 input의 resolution이 변하거나 첫번째 block의 경우 input과 output의 channel이 다르기 때문에 이를 convolution을 통해 다시 맞춰줘야 한다는 것이다. 마지막으로 skip connection은 단순 덧셈으로 구현이 가능하다.

 

BottleNeck Block

class ResBottleNeckBlock(Module):
    expansion = 4

    def __init__(self, in_ch, out_ch, stride=1):
        super().__init__()
        expans_ch = out_ch * self.expansion
        self.conv1 = BnConv(in_ch, out_ch, kernel_size=1, padding=0)
        self.conv2 = BnConv(out_ch, out_ch, kernel_size=3,stride=stride)
        self.conv3 = BnConv(out_ch, expans_ch, kernel_size=1, padding=0)

        if stride !=1 or in_ch != expans_ch:
            self.residual = BnConv(in_ch, expans_ch, kernel_size=1, stride=stride, padding=0)
        else:
            self.residual = nn.Sequential()

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        return out + self.residual(x)

BottleNeck Block의 경우, block의 최종 output이 이전과 다르게 output * expansion이 된다. 왜 이렇게 expansion 변수를 선언해 구별하였는지는 추후에 설명한다.

 

BottlenNeck이라는 구조는 흔히 1x1 convolution을 통해 channel을 바꾸어준 후, 일반적인 convolution을 하고 다시 이를 1x1 convolution으로 channel을 바꾸는 형태의 block을 말한다. 이는 주로 2가지 용도로 쓰이는데, 첫번째는 channel을 낮추어서 연산량을  줄이는 것이고 두번째는 expansion을 함으로써 relu와 같은 activation에서 사라지는 정보를 최대한 줄이도록 하는 것이다. 자세한 내용은 mobilenet논문을 참고하기 바란다. 이 논문에서는 out_ch로 channel을 바꾸어 준 후에 convolution을 하고 이를 out_ch * expansion(ex. 4)의 크기로 바꾼다. 

 

여기서도 stride와 in_ch과 최종 output_ch이 같지 않을 때 convolution을 통해 이를 맞추어 준다.

 

Resnet Body

위의 block을 사용해서 body를 만들어주자! 근데 그냥 단순히 이를 다 나열해버리면 layer의 수가 많은 경우 매우 난잡한 코딩이 될 것이다. 따라서 이를 clean and short하게 바꾸어보자!

make_blocks

    def make_blocks(self, out_ch, num_block, stride=1):
        layers = []
        layers.append(self.block(self.tracking_ch, out_ch, stride=stride)) # first layer
        self.tracking_ch = out_ch * self.block.expansion
        for _ in range(num_block - 1):
            layers.append(self.block(self.tracking_ch, out_ch))
        return nn.Sequential(*layers)

layers 리스트와 반복문을 통해서 이와 같이 간단히 만들어진다. 첫번째 block을 제외한 다른 blocks은 input과 output 모두 out_ch * expansion을 가지므로 이를 계산하여 넣어줄 필요가 있다. 여기서 expansion 변수의 용도가 나오는 것이다. 이러한 변수를 선언해 놓지 않으면 함수를 2개 만들어야 하거나 다른 복잡한 방법을 써야하므로 추천하지 않는다. 마지막으로 *layers를 통해서 Sequential 모듈에 리스트의 모든 elements를 넣으면 끝이다.

body

        self.blocks1 = self.make_blocks(first_ch, num_blocks[0], stride=2)
        self.blocks2 = self.make_blocks(first_ch*2, num_blocks[1], stride=1)
        self.blocks3 = self.make_blocks(first_ch*4, num_blocks[2], stride=1)
        self.blocks4 = self.make_blocks(first_ch*8, num_blocks[3], stride=1)
        self.bn = nn.BatchNorm2d(self.tracking_ch)

위와 같이 함수를 정해놓으면 이처럼 매우 간결하고 깔끔한 body를 작성할 수 있다.

 

Header

마지막으로 Header를 작성하면 끝이다. 이는 보통 fc layer로 이루어지고 마지막에 class 수의 Vector로 매핑한다.

        self.gap = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(self.tracking_ch, num_classes)

논문에서는 fc layer를 마지막에 매핑하는 부분 이외엔 사용하지 않았다. fully connected layer는 overfitting문제와 파라메터 개수가 너무 많아진다는 문제를 가지고 있기 때문에 잘 사용되지 않는다. 여기서는 대신 Global Average Pooling을 사용하여 1차원으로 만듦과 동시에 정보를 압축(?)하는 방식을 사용한다. 

 

ResNet Architecture

이제 forward함수를 통해서 전체적인 architecture를 작성해주면 된다. stem과 body, header를 모두 정의해두었으므로 어려운 것은 없다 ㅎㅎ.

        '''stem'''
        x = self.conv1(x)
        x = self.pool1(x)

        '''body'''
        x = self.blocks1(x)
        x = self.blocks2(x)
        x = self.blocks3(x)
        x = self.blocks4(x)
        x = F.relu(self.bn(x))

        '''head'''
        gap = self.gap(x)
        flat = gap.view(gap.size(0), -1)
        out = self.fc(flat)
        return out

**view의 경우, gap에서 나오는 output이 1x1xout_ch의 shape로 나오기 때문에 이를 1차원으로 reshape해주는 역할을 한다.

 

---

이렇게 Resnet을 구현해보았다. 전체 코드는 다음 github에 올려놓았으니 참고하길 바란다. 이후에도 차근차근 하나씩 만들어나갈 예정인데, 이번처럼 자세하게 하나하나 다룰지 중요한 부분만 다룰지 고민중이다.