DenseNet 구현 및 학습

DenseNetwork 구현 및 학습

• 필자는 구글 colab을 통해 학습시켰으며, @tf.function을 사용하기위해 tensorflow 2.0 버젼을 시용하였다.

!pip install tensorflow==2.0.0-beta1

• 사용할 라이브러리 import

import tensorflow as tf
import numpy as np

• 위에서 설치한 tensorflow의 버전이 2.0인지 확인

print(tf.__version__)

• Hyper parameter를 설정

EPOCHS = 10

DenseUnit 구현

class DenseUnit(tf.keras.Model):
    def __init__(self, filter_out, kernel_size):
        super(DenseUnit, self).__init__()
        # batch normalization -> ReLu -> Conv Layer
        # 여기서 ReLu 같은 경우는 변수가 없는 Layer이므로 여기서 굳이 initialize 해주지 않는다. (call쪽에서 사용하면 되므로)
        # Pre-activation 구조는 똑같이 가져가되, concatenate 구조로 만들어주어야함에 주의하자!!

        self.bn = tf.keras.layers.BatchNormalization()
        self.conv = tf.keras.layers.Conv2D(filter_out, kernel_size, padding='same')
        self.concat = tf.keras.layers.Concatenate()

    def call(self, x, training=False, mask=None): # x : (Batch 갯수, Height, width, Channel_in)
        # training 꼭 잊어버리지 말자!!
        h = self.bn(x, training=training)
        h = tf.nn.relu(h)
        h = self.conv(h) # h : (Batch, height, width, filter_output) zero-padding을 했으므로 filter만 바뀜
        return self.concat([x, h]) # (Batch, height, width, (channel_in + filter_output))

DenseLayer 구현

class DenseLayer(tf.keras.Model):
    def __init__(self, num_unit, growth_rate, kernel_size):
        super(DenseLayer, self).__init__()
        self.sequence = list()
        for idx in range(num_unit):
            self.sequence.append(DenseUnit(growth_rate, kernel_size))

    def call(self, x, training=False, mask=None):
        for unit in self.sequence:
            x = unit(x, training=training)
        return x

Transition Layer 구현

• Maxpooling을 해줄 때 필요함.
• 예를 들어, DenseLayer를 사용을 하게 되면, Growth_rate=32, num_unit=8인 경우에는, 32X8 만큼 channel의 수가 급격하게 증가하기 때문에 너무 커질 수 있다. 이럴 경우 이 Transition Layer를 통해 Pooling 전에 channel의 수를 조절해주기 위해 사용되는 것이다.

class TransitionLayer(tf.keras.Model):
    def __init__(self, filters, kernel_size):
        super(TransitionLayer, self).__init__()
        # transition을 할 경우에는 이렇게 convolution을 해서 단순히 filter 개수를 변경만 해준뒤
        # 그 다음에 Maxpooling을 해주는 식으로 구현이 된다.
        self.conv = tf.keras.layers.Conv2D(filters, kernel_size, padding='same')
        self.pool = tf.keras.layers.MaxPool2D()

    def call(self, x, training=False, mask=None):
        # 여기서는 Batch normalization이 없기 때문에 training을 안써줘도 된다.
        x = self.conv(x)
        return self.pool(x)

모델 정의

class DenseNet(tf.keras.Model):
    def __init__(self):
        super(DenseNet, self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(8, (3, 3), padding='same', activation='relu') # 28x28x8

        # num_unit=2, growth_rate=4, kernel_size=(3,3)
        # num_unit은 ResNet과 동일
        self.dl1 = DenseLayer(2, 4, (3, 3)) # 28x28x(8+2*4)
        self.tr1 = TransitionLayer(16, (3, 3)) # 14x14x16

        self.dl2 = DenseLayer(2, 8, (3, 3)) # 14x14x(16 + 2*8)
        self.tr2 = TransitionLayer(32, (3, 3)) # 7x7x32

        self.dl3 = DenseLayer(2, 16, (3, 3)) # 7x7x(32+2*16)

        self.flatten = tf.keras.layers.Flatten()
        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
        self.dense2 = tf.keras.layers.Dense(10, activation='softmax')       

    def call(self, x, training=False, mask=None):
        x = self.conv1(x)

        x = self.dl1(x, training=training)
        x = self.tr1(x)

        x = self.dl2(x, training=training)
        x = self.tr2(x)

        x = self.dl3(x, training=training)

        x = self.flatten(x)
        x = self.dense1(x)
        return self.dense2(x)

학습, 데스트 루프 정의

# Implement training loop
@tf.function
def train_step(model, images, labels, loss_object, optimizer, train_loss, train_accuracy):
    with tf.GradientTape() as tape:
        predictions = model(images, training=True)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)

    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_loss(loss)
    train_accuracy(labels, predictions)

# Implement algorithm test
@tf.function
def test_step(model, images, labels, loss_object, test_loss, test_accuracy):
    predictions = model(images, training=False)

    t_loss = loss_object(labels, predictions)
    test_loss(t_loss)
    test_accuracy(labels, predictions)

데이터셋 준비

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

x_train = x_train[..., tf.newaxis].astype(np.float32)
x_test = x_test[..., tf.newaxis].astype(np.float32)

train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

학습 환경 정의

모델 생성, 손실함수, 최적화 알고리즘, 평가지표 정의

# 모델 생성
model = DenseNet()

# 손실함수 정의 및 최적화 기법 정의
loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

# 평가지표 정의
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

학습 루프 동작

• 어떤 것은 DenseNet이 성능이 더 좋게 나오고 어떤 것은 ResNet이 더 좋게 나오기도 한다. 또한 parameter에 따라 다르다. Resnet은 좀 더 안정적이게 수렴하지만 Densenet은 좀 fluctuate하다는 점을 유의하자.

for epoch in range(EPOCHS):
  for images, labels in train_ds:
    train_step(model, images, labels, loss_object, optimizer, train_loss, train_accuracy)

  for test_images, test_labels in test_ds:
    test_step(model, test_images, test_labels, loss_object, test_loss, test_accuracy)

  template = "Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}"
  print(template.format(epoch+1,
                        train_loss.result(),
                        train_accuracy.result() * 100,
                        test_loss.result(),
                        test_accuracy.result() * 100))

#train_loss.reset_states()
#train_accuracy.reset_states()
#test_loss.reset_states()
#test_accuray.reset_states()