DenseNet Architecturesの改良

今回は、比較的軽量なモデルであるDenseNetを使い、アーキテクチャをkmnist向けに軽量化してみました。論文で提案されているのは以下のモデルで、Kerasにもデフォルトで搭載されているモデルです。


ブロック数の設計には揺らぎがあることを確認できます。そこで、これらのブロック数やハイパーパラメータをOptunaで最適化したらどうなるかを検証してみました。
まず、DenseBlock(1)(2)の数を3/6に減らしてDenseBlock(3)(4)の編成がどうなるのかを確認しました。その結果、

{'growth_rate': 31, 'compression': 0.36752530507527614, 'n_block3': 31, 'n_block4': 15}

という結果が得られました。つまり、DenseBlock(1)(2)が3, 6層あるとき、DenseBlock(3)(4)は31, 15層必要になりました。3層が比較的重い設定になったのか、学習時間はかなり長めに感じました。更に軽量なモデルを作りたい場合、以下のコードのn_1stblockを2以下にすると良いでしょう。
また、想定よりもcompressionが低下しました。論文では0.5になっていたので、レイヤー毎に削減されるパラメータ数が多く、高速化できていることが分かります。学習率は前半を0.01、中盤から後半にかけて段階毎に0.001, 0.0001と減衰するようにしており、これもDenseNetの元の論文に従っています。
計算資源はNvidiaのTesla K80 1台分です。

ソースコードは以下の通りです。

追記、n_1stblockは1~4で検証して、3が最も良かったですが、5以降は計算時間がかかるため行なっていません。optunaでかなりの時間を割いてしまうようです。

#!/usr/bin/env python3
import os
import numpy as np
import pandas as pd
import keras
import optuna
from keras.utils import np_utils
from keras.models import Model, load_model
from keras.layers import Input, Conv2D, BatchNormalization, Activation, GlobalAveragePooling2D, Dense, Concatenate, AveragePooling2D
from keras.callbacks import ModelCheckpoint, CSVLogger
from sklearn.model_selection import train_test_split
from keras.optimizers import Adam
from keras.callbacks import LearningRateScheduler
from keras.callbacks import EarlyStopping


n_1stblock = 3
n_class = 10
img_shape = (28, 28, 1)

X_train_file = np.load('kmnist-train-imgs.npz')
pred_X_file = np.load('kmnist-test-imgs.npz')
y_train_file = np.load('kmnist-train-labels.npz')

X_train = X_train_file['arr_0']
X_pred = pred_X_file['arr_0']
y_train = y_train_file['arr_0']

X_train = np.expand_dims(X_train, axis = 3) / 255.0
X_pred = np.expand_dims(X_pred, axis = 3) / 255.0
y_train = np_utils.to_categorical(y_train, n_class)


def dense_block(x, k, n_block):
  for i in range(n_block):
    main = x
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    #1x1
    x = Conv2D(filters = 64, kernel_size = (1, 1), padding = 'valid')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    #3x3
    x = Conv2D(filters = k, kernel_size = (3, 3), padding = 'same')(x)
    #concatenate
    x = Concatenate()([main, x])

  return x

def transition_layer(inputs, compression):
  n_channel = int(inputs.shape[3])
  filters = int(n_channel * compression)
  x = Conv2D(filters = filters, kernel_size = (1, 1))(inputs)
  outputs = AveragePooling2D(pool_size = (2, 2))(x)

  return outputs

def DenseNet(growth_rate, compression, n_block3, n_block4):
  inputs = Input(shape = (28, 28, 1))

  x = dense_block(inputs, k = growth_rate, n_block = n_1stblock)
  x = transition_layer(x, compression)
  x = dense_block(x, k = growth_rate, n_block = n_1stblock*2)
  x = transition_layer(x, compression)
  x = dense_block(x, k = growth_rate, n_block = n_block3)
  x = transition_layer(x, compression)
  x = dense_block(x, k = growth_rate, n_block = n_block4)

  x = GlobalAveragePooling2D()(x)

  x = Dense(512)(x)
  x = BatchNormalization()(x)
  x = Activation('relu')(x)

  x = Dense(n_class)(x)

  outputs = Activation('softmax')(x)

  return Model(inputs, outputs)

def step_decay(epoch):
    x = 0.01
    if epoch >= 50: x = 0.001
    if epoch >= 75: x = 0.0001
    return x

def objective(trial):
    growth_rate = trial.suggest_int("growth_rate", 8, 36)
    compression = trial.suggest_uniform("compression", 0.35, 0.65)
    n_block3 = trial.suggest_int("n_block3", n_1stblock*3, n_1stblock*11)
    n_block4 = trial.suggest_int("n_block4", n_1stblock*2, n_1stblock*8)
    model = DenseNet(growth_rate, compression, n_block3, n_block4)
    model.compile(optimizer="Adam", loss='categorical_crossentropy', metrics=['accuracy'])
    history = model.fit(X_train, y_train, epochs=10, validation_split=0.15,
                        callbacks=[lr_decay, EarlyStopping()], batch_size=128)
    return -np.amax(history.history['val_acc'])


lr_decay = LearningRateScheduler(step_decay)

study = optuna.create_study(sampler=optuna.samplers.TPESampler())
study.optimize(objective, n_trials=64)
print('best_params')
print(study.best_params)
print('-1 x best_value')
print(-study.best_value)
print('\n --- sorted --- \n')
sorted_best_params = sorted(study.best_params.items(), key=lambda x : x[0])
for i, k in sorted_best_params:
    print(i + ' : ' + str(k))

model = DenseNet(study.best_params["growth_rate"],
                 study.best_params["compression"],
                 study.best_params["n_block3"],
                 study.best_params["n_block4"])

model.compile(optimizer="Adam", loss='categorical_crossentropy', metrics=['accuracy'])
model.summary()

model.fit(X_train, y_train, epochs=100, validation_split=0.15,
          batch_size=128, callbacks=[lr_decay])

y_prob = model.predict(X_pred)
y_pred = y_prob.argmax(axis=-1)
print(y_pred)
ID = [i + 1 for i in range(0, len(X_pred))]
df = pd.DataFrame({"ImageId": ID, "Label": y_pred})
df.to_csv("result_dn3.csv", index=False)