闇のKerasのTPU術：入力が2つあるモデルと、入力が1つのモデルの係数を相互転移させる方法

KerasとTPUの黒魔術です。Kerasで重み共有レイヤーを使い、入力が2つあるモデルを訓練することを考えます。かなり特殊な状況で、訓練は入力が2つで（2つのサンプルを使って）行いたいものの、実際に推論するときは入力が1つで計算したい場合があります。例えば、Siamese Networkのようなサンプル間の距離を考えながら、埋め込み計算をする場合です。

つまり、「入力が2つのモデルから入力が1つのモデルに係数を転移させる」ということが必要となります。この手の黒魔術はCPUやGPUだとうまくいってTPUだと失敗するケースが多いものの、TPUでも成功したのが美味しいポイントでした。

環境：Colab TPU、TensorFlow v1.12.0

こんなイメージ

MLPやCNNでの想定です。この2つではうまくいくことを確認しましたが、RNNでは確認していないので、もしかするとうまくいかないかもしれません。

この図のように、「入力が1つのモデルから、入力が2つのモデルへ」またはその逆の「入力が2つのモデルから、入力が1つのモデルへ」係数の転移（移植）を相互で行いたいと思います。また入力が2つのモデルでは、全ての層の係数が入力に対して共有されているものとします。

「こんな移植できるの？」と思うかもしれませんができます。

そもそもどうやって全レイヤーの重みが共有されるようなネットワークを作る？

重み共有レイヤー

大きなポイントはネットワークの定義方法です。Kerasのドキュメントにもありますが、重み共有レイヤーのよくありがちな実装では、Functional APIを使い、

from keras.layers import Dense, Input

input1 = Input((784,))
input2 = Input((784,))

some_layer = Dense(128, activation="relu")

x1 = some_layer(input1)
x2 = some_layer(input2)

というように、「1回レイヤーの変数を定義して、その後にKerasのテンソルを計算する」というような計算をします。Functional APIを知らないとこれだけで気持ち悪いような感じはしますが、Sequential APIは入門としては良いものの、ほとんど運用上のメリットがないので、ちょっと突っ込んだことするならFunctional APIを使いましょう。

例えばFunctional APIで入力が1つのモデルでは、

x1 = Dense(128, activation="relu")(input1)

のような書き方をします。重み共有レイヤーでなければこれは間違いではありません。しかし、これでもう1個、

x2 = Dense(128, activation="relu")(input2)

というような書き方をすると、x1とx2の重みは共有されません。なぜなら、レイヤーのインスタンスが異なるためだからです。先程のようなレイヤーのインスタンスを1つにするようにしてあげなければいけません。

先程の書き方はネットワークが浅ければOKです。ただ、ネットワークが深くなってしまうと、2倍、3倍のコードを書かなければいけないのでスマートではありません。なので、CNNのような深いモデルに対してはもうちょっと工夫した書き方をする必要があります。

レイヤーを配列で管理してしまう

一つの案としては、レイヤーのインスタンスをリストとして管理してしまうことです。

from keras import layers

layer_list = []
layer_list.append(layers.Conv2D(32, 3, padding="same"))
layer_list.append(layers.BatchNormalization())
layer_list.append(layers.Activation("relu"))
layer_list.append(layers.AveragePooling2D(2))
layer_list.append(layers.Conv2D(64, 3, padding="same"))
layer_list.append(layers.BatchNormalization())
layer_list.append(layers.Activation("relu"))
layer_list.append(layers.GlobalAveragePooling2D())

もうちょっと深くなったらforループ等で書くのが楽でしょう。このリストに対して、

inputs = [Input((28,28,1)) for i in range(2)]
outputs = []
for input in inputs:
    x = input
    for l in layers:
        x = l(x)
    outputs.append(x)

とforループでレイヤー呼び出しをしてしまいます。モデルの定義は、

from keras.models import Model
model = Model(inputs, outputs)

とすればいいですからね（入力が1個の場合はリストが不要なのでif等で分岐させましょう）。レイヤーのインスタンスをリストに格納し、forループで呼び出すという飛び技を使うことで、重み共有レイヤーがあっても可読性を損なうことがなくなります。入力が1つのモデル、2つのモデルでも全く同じコードから生成することができるのもポイントが高いです。

ちなみに今回は深く掘り下げませんが、ResNetのような分岐構造がある場合はちょっと面倒になります。しかし、AddやConcatのところで区切り、メイン側とSkipConnection側で多重配列にすれば、動的なネットワーク生成はできると思います。VGGやMobileNetのような直線的なモデルの場合は今回の方法でOKです。

入力が1つのモデル→入力が2つのモデルへの係数移植

安直に2つのモデルを作るとうまくいかない

今までの話をまとめて、入力が2つでも1つでも同じように作れる関数を定義しました。

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.contrib.tpu.python.tpu import keras_support
import tensorflow.keras.backend as K

import numpy as np
import os

def create_model(inputs):
    # layers
    layer_list = []
    layer_list.append(layers.Dense(64, activation="relu"))
    layer_list.append(layers.Dense(5, activation="relu"))

    # outputs
    outputs = []
    for x in inputs:
        for l in layer_list:
            x = l(x)
        outputs.append(x)
    if len(outputs) == 1:
        return Model(inputs[0], outputs[0])
    else:
        # 出力層が複数あるとpredictでうまくいかないのでConcatする
        x = layers.Concatenate()(outputs)
        return Model(inputs, x)

ただし入力が2つあるケースで、出力も2つにするとpredictがうまくいかない（shapeが違うぞと言われる）ので、暫定的にConcatenateさせています。これだとエラーは出ません。

そして、入力が1つのモデルと入力が2つのモデルを別々に生成してみます。

係数の同期が取れない例

if __name__ == "__main__":
    K.clear_session()
    X = np.arange(800).reshape(8,100)

    # 別々にモデルを作ると別な出力になる
    # 入力が1つのモデル
    input = layers.Input((100,))
    single_model = create_model([input])
    single_model.compile(tf.train.GradientDescentOptimizer(0.1), "mse") # TPUのモデル変換のためなんでなんでもよい
    single_model.summary()

    # TPUモデルに変換
    tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
    tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
    strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
    single_model = tf.contrib.tpu.keras_to_tpu_model(single_model, strategy=strategy)

    # singleの推論
    y_single = single_model.predict(X)
    print(y_single)

    # 入力が2つのモデル
    input1 = layers.Input((100,))
    input2 = layers.Input((100,))
    double_model = create_model([input1, input2])
    double_model.compile(tf.train.GradientDescentOptimizer(0.1), "mse") # TPUのモデル変換のためなんでなんでもよい
    double_model.summary()

    # TPUモデルに変換
    double_model = tf.contrib.tpu.keras_to_tpu_model(double_model, strategy=strategy)

    # doubleの推論
    y_double = double_model.predict([X, X])
    print(y_double)

これは当然別々の出力を返します（具体的な値は初期値の乱数挙動によって実行するたびに変わります）。レイヤーのインスタンスが違うので当たり前ですね。

出力

# 入力が1つのモデル
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         (None, 100)               0         
_________________________________________________________________
dense (Dense)                (None, 64)                6464      
_________________________________________________________________
dense_1 (Dense)              (None, 5)                 325       
=================================================================
Total params: 6,789
Trainable params: 6,789
Non-trainable params: 0
_________________________________________________________________

# その推論
[[  0.        0.       57.22351   0.        0.     ]
 [  0.        0.       87.0583    0.        0.     ]
 [  0.        0.      115.63985   0.        0.     ]
 [  0.        0.      145.86166   0.        0.     ]
 [  0.        0.      176.26633   0.        0.     ]
 [  0.        0.      209.57587   0.        0.     ]
 [  0.        0.      242.88531   0.        0.     ]
 [  0.        0.      276.19476   0.        0.     ]]

# 入力が2つのモデル
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_2 (InputLayer)            (None, 100)          0                                            
__________________________________________________________________________________________________
input_3 (InputLayer)            (None, 100)          0                                            
__________________________________________________________________________________________________
dense_2 (Dense)                 (None, 64)           6464        input_2[0][0]                    
                                                                 input_3[0][0]                    
__________________________________________________________________________________________________
dense_3 (Dense)                 (None, 5)            325         dense_2[0][0]                    
                                                                 dense_2[1][0]                    
__________________________________________________________________________________________________
concatenate (Concatenate)       (None, 10)           0           dense_3[0][0]                    
                                                                 dense_3[1][0]                    
==================================================================================================
Total params: 6,789
Trainable params: 6,789
Non-trainable params: 0
__________________________________________________________________________________________________

# その推論（columnの0～4と5～9は入力の1つ目、2つ目に対応しているが、内部での重みが共有されているので、同一の値を返す）
[[  0.        38.57988    0.        95.709076  37.404163   0.
   38.57988    0.        95.709076  37.404163]
 [  0.       152.1218     0.       209.1978   140.52289    0.
  152.1218     0.       209.1978   140.52289 ]
 [  0.       257.20923    0.       312.2815   238.54837    0.
  257.20923    0.       312.2815   238.54837 ]
 [  0.       359.7143     0.       412.8396   332.09595    0.
  359.7143     0.       412.8396   332.09595 ]
 [  0.       460.10046    0.       512.0321   419.4279     0.
  460.10046    0.       512.0321   419.4279  ]
 [  0.       560.48676    0.       611.2248   506.75964    0.
  560.48676    0.       611.2248   506.75964 ]
 [  0.       660.87305    0.       710.41736  594.09155    0.
  660.87305    0.       710.41736  594.09155 ]
 [  0.       761.2593     0.       809.61     681.4234     0.
  761.2593     0.       809.61     681.4234  ]]

今、「重み共有レイヤー」と「モデル同士のレイヤー（重み）の共有」を同時に論じているので話が見えにくくなっていますが、入力が2つのモデルでは内部での重みは共有されているものの、入力が2つのモデルと1つのモデルの同期は取れていません。この記事での目的は後者の同期を図るものです。

save_weights, load_weightsで係数を転移させる

実は、（少なくとも同一モデル内の全レイヤーで重みを共有していれば）入力の個数が異なっていても、save_weights()→load_weights()で係数を転移させることができます。これによって、レイヤーのインスタンスが別々でも同一の係数を再現させることが可能になります。

先程の例を少し変えます。

成功例

if __name__ == "__main__":
    K.clear_session()
    X = np.arange(800).reshape(8,100)

    # save_weightsをはさむと入力の個数が違っても同期できる
    # 入力が1つのモデル
    input = layers.Input((100,))
    single_model = create_model([input])    
    single_model.compile(tf.train.GradientDescentOptimizer(0.1), "mse") # TPUのモデル変換のためなんでなんでもよい
    single_model.save_weights("weights.hdf5", save_format="h5")

    # TPUモデルに変換
    tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
    tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
    strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
    single_model = tf.contrib.tpu.keras_to_tpu_model(single_model, strategy=strategy)

    # singleの推論
    y_single = single_model.predict(X)
    print(y_single)

    # 入力が2つのモデルの係数をload_weightsで、入力が1つのモデルから転移
    input1 = layers.Input((100,))
    input2 = layers.Input((100,))
    double_model = create_model([input1, input2])
    double_model.load_weights("weights.hdf5")
    double_model.compile(tf.train.GradientDescentOptimizer(0.1), "mse") # TPUのモデル変換のためなんでなんでもよい

    # TPUモデルに変換
    double_model = tf.contrib.tpu.keras_to_tpu_model(double_model, strategy=strategy)

    # doubleの推論
    y_double = double_model.predict([X, X])
    print(y_double)

    # 結果の確認
    print("double_modelの重み共有 :", np.all(y_double[:,:5]==y_double[:,5:]))
    print("double_modelとsingle_modelの重み共有 :", np.all(y_double[:,:5]==y_single))

本当にうまくいくでしょうか？

# 入力が1つのモデル
[[1.17118584e+02 1.67248783e+01 0.00000000e+00 3.48740425e+01
  0.00000000e+00]
 [3.40153900e+02 4.46984100e+00 0.00000000e+00 5.20611725e+01
  0.00000000e+00]
 [5.67267822e+02 5.66595459e+00 0.00000000e+00 6.38724823e+01
  1.02516174e+00]
 [7.89560669e+02 1.10407867e+01 0.00000000e+00 7.61143646e+01
  3.85089111e+00]
 [1.01189038e+03 1.56332855e+01 0.00000000e+00 8.88464508e+01
  6.88870239e+00]
 [1.23421997e+03 2.02258453e+01 0.00000000e+00 1.01578629e+02
  9.92658997e+00]
 [1.45522437e+03 2.51077271e+01 0.00000000e+00 1.14337677e+02
  1.42870178e+01]
 [1.67615308e+03 3.01563110e+01 0.00000000e+00 1.27196930e+02
  1.91447449e+01]]

# 入力が2つのモデル
[[1.17118584e+02 1.67248783e+01 0.00000000e+00 3.48740425e+01
  0.00000000e+00 1.17118584e+02 1.67248783e+01 0.00000000e+00
  3.48740425e+01 0.00000000e+00]
 [3.40153900e+02 4.46984100e+00 0.00000000e+00 5.20611725e+01
  0.00000000e+00 3.40153900e+02 4.46984100e+00 0.00000000e+00
  5.20611725e+01 0.00000000e+00]
 [5.67267822e+02 5.66595459e+00 0.00000000e+00 6.38724823e+01
  1.02516174e+00 5.67267822e+02 5.66595459e+00 0.00000000e+00
  6.38724823e+01 1.02516174e+00]
 [7.89560669e+02 1.10407867e+01 0.00000000e+00 7.61143646e+01
  3.85089111e+00 7.89560669e+02 1.10407867e+01 0.00000000e+00
  7.61143646e+01 3.85089111e+00]
 [1.01189038e+03 1.56332855e+01 0.00000000e+00 8.88464508e+01
  6.88870239e+00 1.01189038e+03 1.56332855e+01 0.00000000e+00
  8.88464508e+01 6.88870239e+00]
 [1.23421997e+03 2.02258453e+01 0.00000000e+00 1.01578629e+02
  9.92658997e+00 1.23421997e+03 2.02258453e+01 0.00000000e+00
  1.01578629e+02 9.92658997e+00]
 [1.45522437e+03 2.51077271e+01 0.00000000e+00 1.14337677e+02
  1.42870178e+01 1.45522437e+03 2.51077271e+01 0.00000000e+00
  1.14337677e+02 1.42870178e+01]
 [1.67615308e+03 3.01563110e+01 0.00000000e+00 1.27196930e+02
  1.91447449e+01 1.67615308e+03 3.01563110e+01 0.00000000e+00
  1.27196930e+02 1.91447449e+01]]
double_modelの重み共有 : True
double_modelとsingle_modelの重み共有 : True

うまくいきました。設計上、入力が2つのモデル内での重み共有とが行われているのと同時に、入力が1つのモデルと入力が2つのモデルの係数を同期させることに成功しました。入力の個数が違ってもweightsのファイルが適用できてしまうの面白いですね。

入力が2つのモデル→入力が1つのモデルへの係数移植

今度は逆の移植をしてみます。同じようにできます。

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.contrib.tpu.python.tpu import keras_support
import tensorflow.keras.backend as K

import numpy as np
import os

def create_model(inputs):
    # layers
    layer_list = []
    layer_list.append(layers.Dense(64, activation="relu"))
    layer_list.append(layers.Dense(5, activation="relu"))

    # outputs
    outputs = []
    for x in inputs:
        for l in layer_list:
            x = l(x)
        outputs.append(x)
    if len(outputs) == 1:
        return Model(inputs[0], outputs[0])
    else:
        # 出力層が複数あるとpredictでうまくいかないのでConcatする
        x = layers.Concatenate()(outputs)
        return Model(inputs, x)

if __name__ == "__main__":
    K.clear_session()
    X = np.arange(800).reshape(8,100)

    # save_weightsをはさむと入力の個数が違っても同期できる
    # 入力が2つのモデル
    input1 = layers.Input((100,))
    input2 = layers.Input((100,))
    double_model = create_model([input1, input2])
    double_model.compile(tf.train.GradientDescentOptimizer(0.1), "mse") # TPUのモデル変換のためなんでなんでもよい
    double_model.save_weights("weights.hdf5", save_format="h5")

    # TPUモデルに変換
    tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
    tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
    strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
    double_model = tf.contrib.tpu.keras_to_tpu_model(double_model, strategy=strategy)

    # doubleの推論
    y_double = double_model.predict([X, X])
    print(y_double)

    # 入力が1つのモデルの係数をload_weightsで、入力が2つのモデルから転移
    input = layers.Input((100,))
    single_model = create_model([input])
    single_model.load_weights("weights.hdf5")
    single_model.compile(tf.train.GradientDescentOptimizer(0.1), "mse") # TPUのモデル変換のためなんでなんでもよい

    # TPUモデルに変換
    single_model = tf.contrib.tpu.keras_to_tpu_model(single_model, strategy=strategy)

    # singleの推論
    y_single = single_model.predict(X)
    print(y_single)

    # 結果の確認
    print("double_modelの重み共有 :", np.all(y_double[:,:5]==y_double[:,5:]))
    print("double_modelとsingle_modelの重み共有 :", np.all(y_double[:,:5]==y_single))

結果は以下の通りです。

# 入力が2つのモデルの推論
[[ 113.60694    0.         0.         0.         0.       113.60694
     0.         0.         0.         0.     ]
 [ 316.3941     0.         0.         0.         0.       316.3941
     0.         0.         0.         0.     ]
 [ 524.7592     0.         0.         0.         0.       524.7592
     0.         0.         0.         0.     ]
 [ 733.12445    0.         0.         0.         0.       733.12445
     0.         0.         0.         0.     ]
 [ 942.0818     0.         0.         0.         0.       942.0818
     0.         0.         0.         0.     ]
 [1151.6317     0.         0.         0.         0.      1151.6317
     0.         0.         0.         0.     ]
 [1361.1812     0.         0.         0.         0.      1361.1812
     0.         0.         0.         0.     ]
 [1570.731      0.         0.         0.         0.      1570.731
     0.         0.         0.         0.     ]]

# 入力が1つのモデルの推論
[[ 113.60694    0.         0.         0.         0.     ]
 [ 316.3941     0.         0.         0.         0.     ]
 [ 524.7592     0.         0.         0.         0.     ]
 [ 733.12445    0.         0.         0.         0.     ]
 [ 942.0818     0.         0.         0.         0.     ]
 [1151.6317     0.         0.         0.         0.     ]
 [1361.1812     0.         0.         0.         0.     ]
 [1570.731      0.         0.         0.         0.     ]]
double_modelの重み共有 : True
double_modelとsingle_modelの重み共有 : True

load_weights(), save_weights()はシリアライザー感覚で使えるというわけです。

まとめ

• モデル内での重み共有については、Functional APIを使い重み共有レイヤーを使う。この際、例えばレイヤーのインスタンスをリスト等で格納しておいて、動的にネットワークを生成するとスマートに定義できる。
• モデル同士の重みの共有については、（少なくともCNNやMLPのようなモデルで、全レイヤーが重み共有されている状態では）入力の数が変わってもload_weights(), save_weights()で重みの同期を取ることができる
• 係数の保存が、ファイルを通したシリアライザーのような挙動をするので、（少なくとも条件を限定すれば）係数の相互転移ができる
• 別の言い方をすれば、ImageNetで訓練した入力が1つのモデルに対して、学習積み係数をそのままスライドさせて入力を2つや3つにできる。冒頭の例に出した、Siamese Networkを分類モデルから転移学習させることも理論的にはできる

ということでした。以上、KerasとTPUを使った闇の魔術の紹介でした。