[읽고 구현하기] ResNet 완전 처음부터 하나하나 구현하기 - tensorflow, subclassing 방법 (ResNet34)

이번에는 지난 번에 읽었던 ResNet 모델을 직접 한땀한땀 구현해보고자 한다.

리뷰는 여기에, 완성된 전체 코드는 여기에 있다.

 

Tensorflow로 구현했는데, subclassing 방법으로 구현했으니 torch에서도 쉽게 따라할 수 있을 것 같다. 참고로 ResNet34를 처음부터 끝까지 구현해볼 예정이다.

 

1. Basic Block

ResNet의 가장 기본적인 구조, basic block은 residual learning block이다. 아래 그림에서 보면 input -> layer 1 -> relu -> layer 2 -> shortcut connection (F(x)+x) -> relu 로 이루어져 있다. 이 부분만 구현하면 끝! 

 

논문을 좀 더 보다보면, 실제로 논문에서 사용한 구체적인 layer 정보를 볼 수 있다. kernel size 3*3인 convolution layer를 이용했고, filter size는 64이다. 참고로 이미지에는 보이지 않지만 convolution layer바로 다음에는 batch normalization이 항상 따라온다.

전체 코드는 아래와 같다.

# 필요한 module 불러오기
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation, Add
from tensorflow.keras import Model

class ResBlock(Model):
    def __init__(self):
        super(ResBlock, self).__init__()
        # layer 1
        self.conv1 = Conv2D(64, (3,3), padding='same')
        self.bn1 = BatchNormalization()
        self.act1 = Activation('relu')
        
        # layer 2
        self.conv2 = Conv2D(64, (3,3), padding='same')
        self.bn2 = BatchNormalization()
        
        # shortcut connection
        self.add = Add()
        self.act2 = Activation('relu')
    
    def call(self, input):
        # layer 1
        x = self.conv1(input)
        x = self.bn1(x)
        x = self.act1(x)
        
        # layer 2
        x = self.conv2(x)
        x = self.bn2(x)
        
        # shortcut connection
        x = self.add([x, input])
        x = self.act2(x)
        
        return x

 

1. 우선 필요한 layer를 불러오고, class로 ResBlock을 정의해준다.

2. 그리고 첫번째 layer (# layer 1) 부분을 convolution -> batch normalization -> relu 순서대로 정의한다. 

3. 두번째 layer (# layer 2) 는 convolution -> batch normalization 까지만 정의한다. shortcut connection 이후 activation function이 추가되기 때문.

4. shortcut connection은 F(x)+x, 즉 layer를 통과한 결과와 원래 input의 더하기 연산이므로 'Add' layer를 정의한다.

5. 마지막으로 activation function을 덧붙여 준다.

6. 주의: `call` function의 `x = self.add([x, input])` 에서도 볼 수 있듯이,  원래 input과 합쳐주는 과정이 필요하므로, input은 input으로 받고, layer를 통과한 결과를 x로 받는 것이 좋다. 즉 어떻게든 input을 남겨두어야 한다 ㅎㅎ

 

 

미래를 위해 kernel size와 filter 수를 변수로 받을 수 있게 하는 것이 좋다. 그 결과는 다음과 같다.

class ResBlock(Model):
    def __init__(self, filter, kernel_size): #### 다른부분
        super(ResBlock, self).__init__()
        # layer 1
        self.conv1 = Conv2D(filter, kernel_size, padding='same') #### 다른부분
        self.bn1 = BatchNormalization()
        self.act1 = Activation('relu')
        
        # layer 2
        self.conv2 = Conv2D(filter, kernel_size, padding='same') #### 다른부분
        self.bn2 = BatchNormalization()
        
        # shortcut connection
        self.add = Add()
        self.act2 = Activation('relu')
    
    def call(self, input):
        # layer 1
        x = self.conv1(input)
        x = self.bn1(x)
        x = self.act1(x)
        
        # layer 2
        x = self.conv2(x)
        x = self.bn2(x)
        
        # shortcut connection
        x = self.add([x, input])
        x = self.act2(x)
        
        return x

2. Stacked Layer, 즉 residual block의 뭉치(?) 구현하기

다시 우리의 목표인 ResNet34의 구조를 살펴보자. 아래의 왼쪽 그림을 보면, 우선 loop로 이어져 있는 2개의 box 단위가 우리가 위에 구현한 ResBlock에 해당한다. 그런데 이 basic block이 여러개가 쌓여서 하나의 뭉치를 이루는 걸 볼 수 있다. 보라색은 3세트, 연두색은 4세트 등등... 이건 오른쪽의 표(빨간색 박스)에 정리되어 있다. 이번에는 이 뭉치 class를 정의해보려 한다. 

 

각 뭉치들은 filter 혹은 channel 개수로 묶인다. 즉 보라색 뭉치는 filter가 64인 ResBlock, 연두색은 128...

여기서 조금 복잡해지는 게 있다.

 

1. 뭉치가 바뀔 때 처음에 stride=2로 설정하여 output size가 반으로 줄어든다는 것(표 좌측 output size 참고)

2. filter수가 바뀐다는 사실 그 자체이다. input이 filter 64개짜리였는데, ResBlock 끝 부분에서 layer들을 통과한 결과는 filter 128개짜리인 것.

 

ResBlock 마지막에 'Add'를 통해 더해주어야 하는데 반으로 줄어든 output size도 맞지 않고, filter수도 맞지 않다. dimension이 맞지 않아 이대로는 더할 수 없다! 

 

1의 경우는 input을 max pooling 하여 output size를 맞춰줄 수 있다. 그리고 2의 경우에는 아래에서 확인할 수 있듯이 1) zero padding을 filter 단위로 추가해주거나 2) input에 1*1 convolution을 적용하여 개수를 바꿔주는 방법이 있다. 여기서는 1)을 적용해보고자 한다. 

2-1.  Output size decrease, Dimension increase 구현 

zero padding을 이용해 dimension 을 증가시키는 방법은 아래와 같다.

from tensorflow.keras.layers import Lambda, MaxPooling2D

# (1) add zero padding to filter axis
def zeropad(x):
    y =  tf.zeros_like(x)
    return tf.concat([x, y], axis=3)

# (2) get the target filter number
def zeropad_output_shape(input_shape):
    shape = list(input_shape)
    assert len(shape) == 4
    shape[3] *= 2
    return tuple(shape)
    
# (3) decrease output size and increase dimension
def make_shortcut(x, increase_dim):
    if increase_dim:
        x = MaxPooling2D(pool_size=(2,2))(x)
        x = Lambda(zeropad, output_shape=zeropad_output_shape)(x)
    return x

우선 (3)을 보면, 먼저 MaxPooling2D를 이용해 output size 를 반으로 줄여주고, Lambda함수로 zero padding을 하여 filter 수를 2배로 맞춰준다. Lambda는 임의의 표현을 받아서 `Layer` object로 만들어주는 layer 이다. 첫번째 argument로 `Layer` object로 만드는 데에 수행될 `function`이 들어가고, output_shape argument에 `function`으로부터 기대/예상되는 output shape을 받는다. 참고로 `function`에 들어가는 함수의 첫번째 인자는 input tensor이며, `output_shape`에 구체적인 tuple이 아니라 함수가 들어갈 경우 input shape을 인자로 받는다고 가정한다. 그래서 (3)과 같은 함수가 나오는 것! 참고로, `increase_dim`을 받도록 구현했는데 이는 뒤에서 이유를 더 알아볼 수 있다 ㅎㅎ 

 

(1) 은 zero padding 함수, (2)는 input shape으로 부터 filter 개수를 받아서 2배로 만들어준 다음 output shape으로 내뱉는 함수이다.

 

2-2. ResBlock 함수 변경하기 - input filter와 output filter가 다른 경우 dimension 변형

이제 자동으로 input의 shape을 바꿔줄 수 있도록 ResBlock을 조금 더 복잡하게 바꿔보자.

class ResBlock(Model):
    def __init__(self, in_planes, out_planes, kernel_size): #### 다른부분
        super(ResBlock, self).__init__()
        self.increase_dim = True if in_planes != out_planes else False #### 다른부분
        stride = 2 if self.increase_dim else 1 #### 다른부분

        # layer 1
        self.conv1 = Conv2D(out_planes, kernel_size, strides=stride, padding='same') #### 다른부분
        self.bn1 = BatchNormalization()
        self.act1 = Activation('relu')
        
        # layer 2
        self.conv2 = Conv2D(out_planes, kernel_size, padding='same') #### 다른부분
        self.bn2 = BatchNormalization()
        
        # shortcut connection
        self.add = Add()
        self.act2 = Activation('relu')
    
    def call(self, input):
        shortcut = make_shortcut(input, self.increase_dim) #### 다른부분
        # layer 1
        x = self.conv1(input)
        x = self.bn1(x)
        x = self.act1(x)
        
        # layer 2
        x = self.conv2(x)
        x = self.bn2(x)
        # shortcut connection
        x = self.add([x, shortcut])
        x = self.act2(x)
        return x

우선 `in_planes`와 `out_planes`를 통해 input filter와 output filter를 받을 수 있게 했다. 그리고 line 4에서 input filter와 output filter가 다르면 `self.increase_dim`이 True가 되도록 하여 자동으로 dimension resizing을 할 수 있도록 준비했다. 또한 filter 수가 다른 경우에는 첫번째 layer에서 stride가 2가 된다는 뜻이므로, 이 부분도 정의해 주었다. 그 아래 부분은 '# layer 1'에서 strides를 인자로 받는다는 것, `out_planes`를 filter수로 받는다는 것 외에는 동일하다! 

 

또한 `call` 함수에서 위에서 정의한 `make_shortcut`함수를 우선 적용해준다. `self.increase_dim`을 통해 변형 여부를 결정해주므로 여기서는 모든 input에 대해 `make_shortcut`함수를 적용하고, 뒤에서 add를 통해 합쳐준다.

 

2-3. Stacked Layer 정의하기

거의 다 왔다. 이제 뭉치들을 정의해보자. 

class StackedLayer(Model):
    def __init__(self, in_planes, out_planes, num_blocks, kernel_size):
        super(StackedLayer, self).__init__()
        self.resblock = []
        for i in range(num_blocks):
            if i == 0 and in_planes != out_planes:
                self.resblock.append(ResBlock(in_planes, out_planes, kernel_size))
            else:
                self.resblock.append(ResBlock(out_planes, out_planes, kernel_size))
    def call(self, x):
        for i in range(len(self.resblock)):
            x = self.resblock[i](x)
        return x

`in_planes`와 `out_planes`를 통해 input filter 수와 output filter 수를 받고, `num_blocks`를 통해 해당 뭉치에 ResBlock이 몇개가 들어갈지 받아준다. 논문과 위의 표를 확인해보면 64짜리 ResBlock 3개, 128짜리 ResBlock 4개, 256짜리 6개, 512짜리 3개를 각각 쌓는 것을 확인할 수 있다. 마지막으로 kernel_size를 받는다. 다 3*3이지만 그래도 ㅎ ㅎ 

 

그리고 뭉치마다 block 갯수가 다른 것은 for loop로 구현해주었다. self.resblock을 빈 list로 정의한 후, num_block으로 loop를 돌아 그 개수만큼 layer가 정의될 수 있게 했다. loop에서 하나 추가된 조건문은, block의 첫번째이면서 input filter수와 output filter수가 다른 경우 dimension 변형이 일어날 수 있도록 첫번째 인자로 input filter수를 받을 수 있게 한 것이고, 그렇지 않은 경우 첫번째 인자로 output filter 수를 넣어 dimension이 같음을 알려주려는 것이다. 

 

마지막으로 `call`함수에서도 for loop를 이용해주었다. 

 

3. 최종 ResNet34 만들기

다시 아래를 보자. 이제 전체를 보자 input ->  7*7 convolution, filter 64, stride 2 layer -> max pool ->  StackLayer 4개 -> average pool -> 1000-d fc -> softmax 의 구조로 되어있다. StackedLayer까지 다 구현했으니 다른 거는 식은 죽 먹기! 

from tensorflow.keras.layers import GlobalAveragePooling2D, Dense

class ResNet34(Model):
    def __init__(self):
        super(ResNet34, self).__init__()
        self.conv1 = Conv2D(64, (7, 7), strides=2, padding='same')
        self.bn1 = BatchNormalization()
        self.act1 = Activation('relu')
        
        self.pool1 = MaxPooling2D(pool_size=(3,3), strides=2, padding='same')
        
        self.stack1 = StackedLayer(64, 64, 3, (3,3))
        self.stack2 = StackedLayer(64, 128, 4, (3,3))
        self.stack3 = StackedLayer(128, 256, 6, (3,3))
        self.stack4 = StackedLayer(128, 512, 3, (3,3))
        
        self.pool2 = GlobalAveragePooling2D()
        self.dense = Dense(1000)
        self.act2 = Activation('softmax')
        
    def call(self, input):
        x = self.conv1(input)
        x = self.bn1(x)
        x = self.act1(x)
        x = self.pool1(x)
        x = self.stack1(x)
        x = self.stack2(x)
        x = self.stack3(x)
        x = self.stack4(x)
        
        x = self.pool2(x)
        x = self.dense(x)
        x = self.act2(x)
        return x

conv1, bn1, act1을 통해 처음의 layer를 구현해줬고, poo1로 pooling 구현, 이후 stack1, 2, 3, 4를 각각 입력해줬다. 마지막으로 global average pooling, dense layer, activation까지 해주면 진짜 진짜 완료!!! 

 

전체가 합쳐진 최종본과 bottleneck block이 사용된 ResNet50은 github에 완결된 코드로 업로드했다! 링크는 여기