ColabのTPUを使った高速なスタイル変換を作った（Neural Style Transfer） #Python

「ディープラーニングの画像生成で新しい千円札作りたくない？？　作りたくない？？」

「Neural Style Trasnferって遅い、コード面倒くさい、L-BFGSってこれ以外使わねえ」→結論：めんどくさいって思ったので、ColabのTPUに移植しました。結構速く動いて、256×256の解像度で3000エポック回して15分で終わります。

元の論文

A Neural Algorithm of Artistic Style : https://arxiv.org/abs/1508.06576
Image Style Transfer Using Convolutional Neural Networks : https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Gatys_Image_Style_Transfer_CVPR_2016_paper.pdf
Perceptual Losses for Real-Time Style Transfer and Super-Resolution : https://arxiv.org/abs/1603.08155

コード

全体のコードはこちら参照
https://github.com/koshian2/NeuralStyleTransferTPU

結果

全てのケースでContent Lossの係数は0.025、Style Lossの係数は1に固定しました。Total Variation Lossのみ変化させています。

新・千円札

北里柴三郎＋葛飾北斎の「神奈川沖浪裏」。画像はリンク元より。

コンテンツ画像

スタイル画像

生成画像（TV=0）

Total Variation Lossがないケース。元の顔写真は維持しつつも、波打っている感じになっています。

生成画像（TV=1e-5）

気持ちTotal Variation Lossを入れたケース。細かいラインは消えてはいるものの、波という感じになっています。ハンコっぽいような印象になっています。

清水寺＋鵺

こっちのほうが面白い結果になっています。清水寺の画像に歌川国芳の鵺（木曽街道六十九次）をスタイルとして入れてみました。

コンテンツ画像

スタイル画像

生成画像（TV=0.01）

清水寺の輪郭は維持しつつも、鵺の輪郭がところどころにあらわれているのがわかるでしょうか。木のラインと空のラインにそれが見られますね。

生成画像（TV=0.25）

少しTotal Variation Lossを増やして平滑化します。清水寺の細かな輪郭は消えましたが、鵺の禍々しい雰囲気が増してまさに「丑三つ時」という感じになっています。

方針

KerasのNeural Style Transferのサンプルを参考にしたが、TPUに移植する過程でもはやコードが別物になった
KerasのVGGはテンソル指定して初期化するとハマる（load_weightsの箇所で「by_name=True」がない）ので、VGG19のコードを自分で作る。KerasのGitHubをコピペし、by_name=Trueを入れて、複数モデル展開してもレイヤーの名前衝突が起きないように改造すればOK。
関連issue：https://github.com/keras-team/keras/issues/4465
Neural Style Transferには、Content Loss、Style Loss、Total Variation Lossの3種類の損失関数がある。Total Variation Lossを入れないのがオリジナルの論文だが、入れたほうが綺麗に仕上がる（これは別の論文）。
Style Lossの計算には訓練済みVGGの特徴量を使う。一般的なGPUのコードだと、損失値や損失関数内でVGGの初期化をするが、TPUの場合は損失関数内でload_weightsの処理が展開できないので、モデル側に損失値を計算するレイヤーを作る。モデル側に作ればCPUで一旦初期化してからTPUに展開するのでうまくいく（ここらへんの話は技術書典の本の8章で書きました）。
損失値をレイヤー側で事前に計算して、損失関数はその値を引き出すだけ
TPUは2019年4月現在、2つ以上の出力があると損失関数や評価関数の表示がおかしくなるというバグがあるので、出力画像の4チャンネル目に損失値を入れて1つの出力にする
オプティマイザはAdamを使う
生成画像を再帰的に訓練させるには、出力画像を係数に持ったレイヤーを定義する。

出力画像を係数に持ったレイヤーとはこんな感じのもの。

# Savng layer of a generated image
class GeneratedImages(layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def build(self, input_shape):
        self.image_params = self.add_weight(name="image_params",
                                shape=(int(input_shape[1]), int(input_shape[2]), int(input_shape[3])),
                                initializer="uniform", trainable=True)
        super().build(input_shape)

    def call(self, x):
        # Apply sigmoid function to the params
        x = K.sigmoid(self.image_params) * K.sign(K.abs(x)+1)
        # [0,1] -> [0, 255]
        x = 255*x
        # RGB -> BGR
        x = x[:, :, :, ::-1]
        # Convert to caffe color scale
        mean = K.variable(np.array([103.939, 116.779, 123.68], np.float32).reshape(1,1,1,3))
        return x - mean
        #return self.image_params * K.sign(K.abs(x)+1)


    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        return super().get_config()

訓練開始時に値域が[0,1]かつRGBのカラースケールで重みを初期化し、VGGのカラースケールに合うように出力をする（レイヤー内で前処理をする）。シグモイド関数は入れても入れなくてもいいが、入れると結構誤差の収束が早くなる。シグモイド関数は入れないからといって、きれいになるわけではなさそう。→シグモイド入れたほうがよくない？ってことになりました。

シグモイド関数で[0,1]出力すればいいので、初期値として与えるコンテンツ画像は、[0,1]スケールのRGB値をロジット変換したものでOKですね。

$$y=\ln\frac{x}{1-x+\epsilon} $$

xを[0,1]スケールのRGB値、yをロジットとします。シグモイド関数の逆関数なので、これにシグモイド関数をかければ[0,1]スケールに戻ります。こんな感じでセットします。

    # initialize generating layer with content image
    # [0,1] -> logit
    content_r = content_numpy.astype(np.float32) / 255.0
    content_logit = np.log(content_r / (1 - content_r + 1e-8))
    model.get_layer("generated").set_weights([content_logit])

ColabのTPUを使った高速なスタイル変換を作った（Neural Style Transfer）

元の論文

コード

結果

新・千円札

清水寺＋鵺

方針

関連