VGGにおける前処理をネットワーク内で行い自作generatorを高速化する(keras)

今回したこと

VGG16をfinetuneする際に自作generatorがボトルネックになっており，時間がかかっていた．
そこで自作generatorの実装を見直すことで 770s/epoch -> 490s/epoch へと学習自体の高速化を図った．

なお，今回用いたVGG16はkeras標準のものではなく，自分でcaffeモデルから変換したものです．（下記事参照
catdance124.hatenablog.jp
keras.applicationのVGG16はすでにBGR -> RGBへと入力が補正されているらしいので注意．

KerasにおけるVGG16の重みは、Oxford大学のVGGによりCreative Commons Attribution Licenseの下で公開されたものを移植しています。そのため、本来、期待する前処理は、BGR順で0～255の値からImageNetのmeanを引いた値となります。ただし、CaffeModelからKerasModelへの変換の過程でRGB順への補正は行われているようですので、RGB順で0～255がKerasとして期待する入力となります。

blog.abars.biz

今回示すコードのまとめは下記colabに載せてあります．
https://colab.research.google.com/drive/1g_W05-9wSzbU5y-M3pDxkGhfGiH77FR4

ネットワーク内でVGG前処理を行う方法

Lambdaで自作レイヤーを作成する．
vgg_preprocess()が前処理を行う関数

import numpy as np
from keras.models import Model
from keras.layers import Input, Lambda, Dense, Flatten

vgg_mean = np.array([123.68, 116.779, 103.939]).reshape((1,1,1,3))
def vgg_preprocess(x):
    x = x - vgg_mean     # subtract mean
    return x[:, :, :, ::-1]    # bgr->rgb


inputs = Input(shape=(224, 224, 3))
x = Lambda(vgg_preprocess, input_shape=(224, 224, 3), name='VGG_preprocess')(inputs)
preprocess_model= Model(inputs=inputs, outputs=x)
preprocess_model.summary()
# _________________________________________________________________
# Layer (type)                 Output Shape              Param #   
# =================================================================
# input_1 (InputLayer)         (None, 224, 224, 3)       0         
# _________________________________________________________________
# VGG_preprocess (Lambda)      (None, 224, 224, 3)       0         
# =================================================================
# Total params: 0
# Trainable params: 0
# Non-trainable params: 0
# _________________________________________________________________

サンプルとして全画素値200の入力を用意

sample = np.zeros((1,224,224,3))
sample[:] = 200.0
print(sample)
# array([[[[200., 200., 200.],
#          [200., 200., 200.],
#          [200., 200., 200.],
#          ...,

出力を見てみる．
結果から前処理（RGB平均値を引く，RGB->BGR変換）が行われていることがわかる．

pred = model.predict(sample)
print(pred)
# array([[[[96.061, 83.221, 76.32 ],
#          [96.061, 83.221, 76.32 ],
#          [96.061, 83.221, 76.32 ],
#          ...,

これでpreprocessをネットワーク内で行うモデルを作成できた．
自作generator内では前処理をする必要がなくなり，GPUパワーで前処理を行うことで高速化が図れる．（本当にGPUで処理されているかは謎）

次にVGG16の特徴抽出部分を用意する．
（冒頭で自分で用意したvgg16を使うと書いたが，ここでは簡略化のためkeras.applications.vgg16を使用する．）
finetuneなのでinclude_top=Falseで，上記モデルの出力を入力するのでinput_shape=preprocess_model.output_shape[1:]としておく．
v

from keras.applications.vgg16 import VGG16
vgg16 = VGG16(include_top=False, weights='imagenet', input_tensor=None, input_shape=preprocess_model.output_shape[1:])
vgg16.summary()
# _________________________________________________________________
# Layer (type)                 Output Shape              Param #   
# =================================================================
# input_18 (InputLayer)        (None, 224, 224, 3)       0         
# _________________________________________________________________
# block1_conv1 (Conv2D)        (None, 224, 224, 64)      1792      
# _________________________________________________________________
# ...
# _________________________________________________________________
# block5_conv3 (Conv2D)        (None, 14, 14, 512)       2359808   
# _________________________________________________________________
# block5_pool (MaxPooling2D)   (None, 7, 7, 512)         0         
# =================================================================

最後にFC層モデルを用意する．
同じようにshape=vgg16.output_shape[1:]．

inputs = Input(shape=vgg16.output_shape[1:])
x = Flatten(name='flatten')(inputs)
x = Dense(4096, activation='relu', name='FC1')(x)
x = Dense(4096, activation='relu', name='FC2')(x)
x = Dense(30, activation='softmax', name='predictions')(x)
top_model = Model(inputs=inputs, outputs=x, name='top_model')
top_model.summary()
# _________________________________________________________________
# Layer (type)                 Output Shape              Param #   
# =================================================================
# input_19 (InputLayer)        (None, 7, 7, 512)         0         
# _________________________________________________________________
# flatten (Flatten)            (None, 25088)             0         
# _________________________________________________________________
# FC1 (Dense)                  (None, 4096)              102764544 
# _________________________________________________________________
# FC2 (Dense)                  (None, 4096)              16781312  
# _________________________________________________________________
# predictions (Dense)          (None, 30)                122910    
# =================================================================

上記3モデル（前処理モデル，VGG16モデル，FC層モデル）を1つのモデルに結合し，前処理をネットワーク内で行うVGG16が構築された．

model = Model(inputs=preprocess_model.input, outputs=top_model(vgg16(preprocess_model.output)))
model.summary()
# _________________________________________________________________
# Layer (type)                 Output Shape              Param #   
# =================================================================
# input_16 (InputLayer)        (None, 224, 224, 3)       0         
# _________________________________________________________________
# VGG_preprocess (Lambda)      (None, 224, 224, 3)       0         
# _________________________________________________________________
# vgg16 (Model)                (None, 7, 7, 512)         14714688  
# _________________________________________________________________
# top_model (Model)            (None, 30)                119668766 
# =================================================================
# Total params: 134,383,454
# Trainable params: 134,383,454
# Non-trainable params: 0
# _________________________________________________________________