Galapagos Tech Blog

株式会社ガラパゴスのメンバーによる技術ブログです。

Kerasのお勉強のついでにStyleTransfer

これはGalapagos Advent Calendar 19日目の記事です。

こんにちは、AIチームの中の人まんだです。 アドベントカレンダー二度目の登場です。前回は合同勉強会の参加レポートを投稿しましたが、 今回はせめてAIチームらしく、機械学習の話題をしたいと思います。

今回の内容は、「Neural-Style-Transferを題材にKerasのお勉強を始めよう」です。

背景

ガラパゴスAIチームでは長らく機械学習モデルはTensorFlowで書いていたのですが、 最近「Kerasって便利そうだよね〜。」「次に新しく組むモデルからはKerasを採用しても良いのでは?」 という話がちらほらと上がり始めていました。 これは、勉強せないかんなぁ〜と考えていたところで、12月のアドベントカレンダーの話が来たので、 「そうだ!アドベントカレンダーねたに抱き合わせでKerasのお勉強をしてしまおう」と思い立ってこの記事が作られています。 ただKerasのお勉強をしてもつまらないので、何かモデルをいじってみようということで、今回はNeural-Style-Transferを題材に選んでみました。

Neural-Style-Transferとは

NeuralStyleTransferとは、論文1で2016年に発表された、画像(絵画とか)のスタイル(画風・雰囲気)を別の画像(写真など)に転写できるようにするネットワークのことです。

このモデルでは、1回の画像のスタイル変換ごとにノイズ画像を出発地点として、ForwardとBackwardの計算が必要になるために出力されるまでそこそこ時間がかかります。 なので、今回ターゲットにするのは、リアルタイムでスタイル変換することを可能にしたモデル2を使います。

StyleTransferとKeras

というわけで、Style-TranferをKerasで...と思ったら、当然のようにすでに実装はされているわけでして。 今回はこの実装コードを参考にしながらKerasのお勉強を進めていきます。

モデルの中身

高速スタイル変換の中身自体の解説は、この記事などで詳しく紹介されています。 ここでは、kerasでの実装コードを眺めてどのようにモデルが記述されているか見てみましょう。

Kerasでのモデルの定義

Kerasで利用出来るモデルは次の二つがあります。

  1. Sequential Model
  2. functional APIを用いたモデル

Sequential Model

以下はドキュメントから持ってきた例です。

model = Sequential()
model.add(Dense(32, input_shape=(500,)))
model.add(Dense(10, activation='softmax'))
model.compile(optimizer='rmsprop',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

addを使ってレイヤーをただただ追加していくシンプルな作りですね。

functional API

シンプルなモデルだったらSequentialモデルで良いのですが、少し洒落たことがしたいとなったら functional APIを使いましょう。 またまたドキュメントから例を持ってくると

from keras.models import Model
from keras.layers import Input, Dense

a = Input(shape=(32,))
b = Dense(32)(a)
model = Model(inputs=a, outputs=b)

と書けます。

Neural-Style-TransferでのモデルをKerasで書くと?

というわけで、さっそくNeural-Style-TransferのモデルをKerasで定義してみましょう。 Githubにあるtrain.pyの中身を見てみると

net = nets.image_transform_net(img_width, img_height, tv_weight)
model = nets.loss_net(net.output, net.input, img_width, img_height, style_image_path, content_weight, style_weight)
model.summary()
optimizer = Adam()  # Adam(lr=learning_rate,beta_1=0.99)
model.compile(optimizer,  dummy_loss)  # Dummy loss since we are learning from regularizes

どうやらnets.pyにスタイル変換ネットワークとロスネットワークが記述されている様子です。 ともあれmodelが定義されたらoptimizerを指定して、model.compileするだけ。お手軽ですね。 今回別口でlossを定義しているので、compile時点ではdummy_lossを指定してます。

スタイル変換ネットワーク

では、nets.image_transform_net()の中身を確認してみましょう。そっとnets.pyを開いてみます。

from keras.layers import Input
from keras.layers.merge import concatenate
from keras.models import Model, Sequential
from layers import InputNormalize, VGGNormalize, ReflectionPadding2D, Denormalize, conv_bn_relu, res_conv, dconv_bn_nolinear
from loss import StyleReconstructionRegularizer, FeatureReconstructionRegularizer, TVRegularizer
from keras import backend as K
from VGG16 import VGG16
import img_util


def image_transform_net(img_width, img_height, tv_weight=1):
    x = Input(shape=(img_width, img_height, 3))
    a = InputNormalize()(x)
    a = conv_bn_relu(32, 9, 9, stride=(1, 1))(a)
    a = conv_bn_relu(64, 3, 3, stride=(2, 2))(a)
    a = conv_bn_relu(128, 3, 3, stride=(2, 2))(a)
    for i in range(5):
        a = res_conv(128, 3, 3)(a)
    a = dconv_bn_nolinear(64, 3, 3)(a)
    a = dconv_bn_nolinear(32, 3, 3)(a)
    a = dconv_bn_nolinear(3, 9, 9, stride=(1, 1), activation="tanh")(a)
    # Scale output to range [0, 255] via custom Denormalize layer
    y = Denormalize(name='transform_output')(a)

    model = Model(inputs=x, outputs=y)

    if tv_weight > 0:
        add_total_variation_loss(model.layers[-1], tv_weight)

    return model

ふむふむ。keras.layersからインポートされているInputで入力を規定している様子。その後InputNormalizeは自作しているようですが、[0, 255]→[0, 1]に正規化しているご様子。 その後、conv_bn_reluが3回繰り返されています。

from keras.layers.core import Activation
from keras.layers.normalization import BatchNormalization
from keras.layers.convolutional import Conv2D
def conv_bn_relu(nb_filter, nb_row, nb_col, stride):   
    def conv_func(x):
        x = ReflectionPadding2D(padding=(nb_row // 2, nb_col // 2))(x)
        x = Conv2D(nb_filter, (nb_row, nb_col), strides=stride, padding='valid')(x)
        x = BatchNormalization()(x)
        x = Activation("relu")(x)
        return x
    return conv_func

これはConvでたたみ込み→バッチ正規化→Reluの処理をしています。 その後はResidual Blockを5層用意しています。さらにdconv_bn_nolinearでDeconvolutionしていますね。この辺りは中身はconv_bn_reluと大差ないので割愛します。 deconv_bn_nolinear()ではactivation関数はデフォルトがReluなのですが最後の層だけはtanhを指定しています。引数で渡すだけで切り替えられて便利ですね。 最後にDenormalize()で[0, 255]の範囲に戻しています。 そしてinputとoutputをModel()に渡してあげて完了です。

Lossネットワーク

VGG16の重みをセットして、Style LossとContent Lossをそれぞれ計算します。

from VGG16 import VGG16
def loss_net(x_in, trux_x_in, width, height, style_image_path, content_weight, style_weight):
    # Append the initial input to the FastNet input to the VGG inputs
    x = concatenate([x_in, trux_x_in], axis=0)
    
    # Normalize the inputs via custom VGG Normalization layer
    x = VGGNormalize(name="vgg_normalize")(x)

    vgg = VGG16(include_top=False, input_tensor=x)

    vgg_output_dict = dict([(layer.name, layer.output) for layer in vgg.layers[-18:]])
    vgg_layers = dict([(layer.name, layer) for layer in vgg.layers[-18:]])

    if style_weight > 0:
        add_style_loss(vgg, style_image_path, vgg_layers, vgg_output_dict, width, height, style_weight)

    if content_weight > 0:
        add_content_loss(vgg_layers, vgg_output_dict, content_weight)

    # Freeze all VGG layers
    for layer in vgg.layers[-19:]:
        layer.trainable = False

    return vgg

Style LossとContent Lossに関しては、元論文に準じています。詳細は元論文を参照してください。

学習

それではいよいよ学習してみましょう。 用意するものは、

  1. 学習したいスタイル画像1枚
  2. 学習用コンテンツ画像 たくさん(今回はMS COCOを使用)
  3. GPUマシン

今回は学習に社内のGPUマシン(GeForce GTX 1080搭載)を使いました。(社内マシンはDockerで運用しているので、このモデルも諸々Dockerに載せる作業がありましたが、詳細は割愛します。) 学習させるスタイルは、モネの睡蓮にしてみました。

f:id:glpgsinc:20171218140434j:plain (http://www.mam-e.it/wp-content/uploads/2017/01/NINFEE-claude-monet-992x538.jpg)

いざ学習へ

from keras.preprocessing.image import ImageDataGenerator
datagen = ImageDataGenerator()
dummy_y = np.zeros((train_batchsize, img_width, img_height, 3))  # Dummy output, not used since we use regularizers to train
skip_to = 0

i = 0
t1 = time.time()
for x in datagen.flow_from_directory(train_image_path, class_mode=None, batch_size=train_batchsize, target_size=(img_width, img_height), shuffle=False):

    if i > nb_epoch:
        break

    if i < skip_to:
        i += train_batchsize
        if i % 1000 == 0:
            print("skip to: %d" % i)
        continue

    hist = model.train_on_batch(x, dummy_y)
    if i % 50 == 0:
        print(hist, (time.time() - t1))
        t1 = time.time()

    if i % 500 == 0:
        print("epoc: ", i)
        val_x = net.predict(x)

        display_img(i, x[0], style)
        display_img(i, val_x[0], style, True)
        model.save_weights(os.path.join(model_dump_dir, style) + '_weights.h5')
        model.save(os.path.join(model_dump_dir, style) + '.h5', include_optimizer=False)
    i += train_batchsize

学習自体はいたってシンプルです。 ImageDataGenerator().flow_from_directoryでバッチサイズ毎の画像読み出しを行って、model.train_on_batch()を呼ぶだけです。シンプル! 一応500ステップ毎にモデルのダンプをしています。

学習途中の画像はこんな感じになっています。

f:id:glpgsinc:20171218141551p:plainf:id:glpgsinc:20171218141555p:plainf:id:glpgsinc:20171218141602p:plainf:id:glpgsinc:20171218141559p:plainf:id:glpgsinc:20171218141611p:plainf:id:glpgsinc:20171218141607p:plain

最後の段は、約16万ステップでの出力になります。しっかり草原がモネの睡蓮っぽく(?)なってますね! きちんとスタイルが学習されています。よかったよかった。

まとめ

Kerasのお勉強を兼ねて、Neural-Style-Transferのモデルを学習させてみました。 実際のところキチンとKerasを使いこなすには、もう少し自分でちゃんと実装してみないとわかりませんね。 でもかなり便利に使えそうなので、じわじわ使っていこうと思います。

おわりに

明日はiosチームの高橋さんが、アプリ開発会社ガラパゴスらしく、このStyleTransferモデルを スマートフォンに載せてリアルタイム画風変換をやってくれるようです。お楽しみに。

さらにおわりに

弊社では、Kerasでバリバリ機械学習したい方、TensorFlowでガシガシ実装したい方、 機械学習興味ある方を絶賛募集中です。ご応募お待ちしております。

www.glpgs.com

以上になります。明日以降もアドベントカレンダーお楽しみに!


  1. Gatys, Leon A., Alexander S. Ecker, and Matthias Bethge. “Image style transfer using convolutional neural networks.” Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR). 2016.

  2. Justin Johnson, Alexandre Alahi, Li Fei-Fei. “Perceptual Losses for Real-Time Style Transfer and Super-Resolution.“ arXiv. 2016