[ML] MobileNetV2 with TensorFlow

이전 글에서 Depthwise Separable Convolution에 대해 알아봤습니다.

이어서 오늘은 TensorFlow를 활용해 MobileNetV2를 만들어 보겠습니다.

Reference

[1] MobileNetV2 논문

MobileNetV2

MobileNetV2는 저자가 앞서 발표한 MobileNetV1을 개선한 모델입니다.

MobileNet의 목표는 딥러닝 모델을 고성능의 GPU를 갖고있는 플랫폼이 아닌 스마트폰과 같은 모바일, 그리고 임베디드 플랫폼과 같은 엣지 디바이스에서 사용할 수 있는 가벼우면서 성능이 높은 모델로 만드는 것입니다.

MobileNetV2는 기존 depthwise separable convolution에 inverted residual 블록을 추가한 것이 핵심입니다.

Inverted residual block [1]

위 그림처럼 ResNet이 제시한 Wide-Narrow-Wide 구조의 Bottleneck이 아닌 Narrow-Wide-Narrow 구조를 사용합니다. 

이는 저차원일수록 ReLU 연산을 거치면 정보 손실이 크기 때문입니다. 아래 그림은 논문에서 제시한 차원별 manifold입니다. 차원이 클수록 기존 정보를 조금 더 잘 보존하면서 압축함을 볼 수 있습니다.

ReLU transformations of dimensional manifolds [1]

Inverted Residual Block

위 표는 inverted residual block 구조를 나타냅니다. \( (h, w, k) \) shape를 갖는 input을 \( (\frac{h}{s}, \frac{w}{s}, k') \)으로 변환합니다. 여기서 s는 stride를 의미하며 t는 expansion factor로 얼마나 채널을 확장할지를 의미합니다. 

즉, bottleneck block의 중간 Wide 구조에 t만큼 채널이 확장됩니다.

Inverted residual block을 3개 파트로 나눌 수 있을 것입니다.

• Expansion factor를 활용해 t만큼 input channel를 확장해 wide 구조 생성
• Depthwise convolution 적용
• Pointwise convolution 적용

아래 코드는 위 3 파트를 의미합니다.

def expansion_layer(x, t, c_i):
    '''
    Expansion layer
    :param x: Tensor, input data
    :param t: int, expansion factor
    :param c_i: int, num of input blocks channel
    :return: Tensor, expanded layer
    '''
    num_filter = t * c_i
    x = tf.keras.layers.Conv2D(filters=num_filter, kernel_size=(1, 1), padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.ReLU(6)(x)
    return x


def depthwise_layer(x, s):
    '''
    Depthwise convolution layer
    :param x: Tensor, input data
    :param s: int, stride. Could be 1 or 2
    :return: Tensor, depthwised layer
    '''
    x = tf.keras.layers.DepthwiseConv2D(kernel_size=3, strides=(s, s), padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.ReLU(6)(x)
    return x


def projection_layer(x, c_o):
    '''
    Pointwise convolution layer
    :param x: Tensor, input data
    :param c_o: intk num of output blocks channel
    :return: Tensor, end of block
    '''
    num_filter = c_o
    x = tf.keras.layers.Conv2D(num_filter, kernel_size=(1, 1), padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    return x3
    

def bottleneck_block(x, t, c, s):
    '''
    baseline of inverted residual block
    :param x: Tensor, input data
    :param t: int, expansion factor
    :param c: int, same as num of output channel
    :param s: int, stride. Could be 1 or 2
    :return: Tensor
    '''
    assert s in [1, 2]
    c_i = x.shape[-1]
    c_o = c
    inp_block = expansion_layer(x, t, c_i)
    inp_block = depthwise_layer(inp_block, s)
    inp_block = projection_layer(inp_block, c_o)
    if x.shape == inp_block.shape:
        inp_block = tf.keras.layers.add([x, inp_block])
    return inp_block

MobileNetV2

MobileNetV2는 앞서 설명한 inverted residual block을 쌓아 만들어지며 세부 구조는 아래 표와 같습니다.

여기서 n은 block을 몇 번 반복할지를 의미합니다.

앞서 코드에서 t, c, s를 받아서 처리할 수 있도록 구축했으니 MobileNet을 만들 때는 아래 코드처럼 블록을 n번 반복해 쌓는 함수를 사용합니다.

def inverted_residual_block(x, t, c, n, s):
    '''
    Repeating the bottleneck block
    :param x: Tensor, input data
    :param t: int, expansion factor
    :param c: int, num of output channel
    :param n: int, repeated n times
    :param s: int, stride
    :return: Tensor
    '''
    for i in range(n):
        x = bottleneck_block(x, t, c, s)
    return

MobileNetV2 코드

def MobileNetV2(input_shape, num_classes=1000):
    '''
    MobilNetV2
    :param input_shape: tuple consist of 3 integers
    :param num_classes: int. num of labels
    :return: mobilenetv2 model
    '''
    input_ = tf.keras.layers.Input(shape=input_shape)
    x = tf.keras.layers.Conv2D(32, kernel_size=(3, 3), strides=2)(input_)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.ReLU(6)(x)

    x = inverted_residual_block(x, 1, 16, 1, 1)
    x = inverted_residual_block(x, 6, 24, 2, 2)
    x = inverted_residual_block(x, 6, 32, 3, 2)
    x = inverted_residual_block(x, 6, 64, 4, 2)
    x = inverted_residual_block(x, 6, 96, 3, 1)
    x = inverted_residual_block(x, 6, 160, 3, 2)
    x = inverted_residual_block(x, 6, 320, 1, 1)

    x = tf.keras.layers.Conv2D(filters=1280, kernel_size=(1, 1), strides=(1, 1))(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.ReLU(6)(x)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dropout(0.3)(x)
    output_ = tf.keras.layers.Dense(num_classes, activation='softmax')(x)

    model = tf.keras.Model(input_, output_)
    return model

코드는 github에서 확인 가능합니다(Github 주소)

다음 글에서는 Quantization friendly MobileNetV2를 만들어보도록 하겠습니다.