tf.keras 流な weight 宣言の方法と罠

tensorflow 2.0 の紹介（日本語訳）
でも書いたとおり、 tensorflow 2.0 からはより Pythonic な、つまり keras ライクなモデルの作り方が主流になっていくようです。
この記事では tf.keras でのレイヤー・モデルの作り方、その中での weight の宣言のしかたとその罠を紹介します。

先にまとめ

• layer/model を使うと weight の一覧をオブジェクト指向的に取れて良い。
• weight (variable) を直接使うものは Layer で実装
  • ただし内部に他の layer など重みを持ちうるものは持たないほうがいい
• layer を使って Model を作る
  • Model は別の Model を含むこともできる
  • 内部に持つ layer, model は必ず self. に入れておくこと

Layer と Model

tf.keras では Layer と Model を組み合わせて学習対象の Model を作っていきます。

上の図のように weight < Layer < Model という包含関係になっています。
Layer は重み（以前の書き方での tf.Variable）を持つことができます。
Model は内部で Layer や別の Model を組み合わせて新たなモデルを作れます。

これら layer, model は layer.weights, model.weights といった形で簡単に重みの一覧を取得することができ、 pre-train / transfer-learning で便利です。

tf.keras.layers.Layer

Layer は input に対する小さな処理の単位です。 Layer は weight (variable) を持つことができます。
weight を宣言する場合、主に build() の中で self.add_weight(name, shape, ...) を使います。

雛形

参考にしたのは公式の Eager Execution 及び tf.keras.layers.Layer です。
※公式サンプルは tensorflow 1.12 で動かなかったので少し修正しています。

Layer を継承してクラスを作り、以下のメソッドを override します。

• init: config (上の例だと output_units)を self に格納する
  • インスタンス作成時に1回呼ばれる
• build: config と input_shape を使って重みの宣言をする
  • 初めてレイヤーを __call__() したときに1回呼ばれる
• call: build で作成した重みを使って計算グラフを構築する（eager の場合は実行する）
  • __call__ が呼ばれるたびに実行されます。

class MyDenseLayer(tf.keras.layers.Layer):
    def __init__(self, output_units, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.output_units = output_units

    def build(self, input_shape):
        # build では変数の宣言と登録を行います。
        # build は最初にレイヤーが実行（callが呼ばれたとき）の1回のみ実行されます。
        # 変数の宣言と登録は、 __init__ で行っても構いません。
        # ただ、 build で行うと入力の shape がわかるため便利というだけです。
        depth = int(input_shape[-1])
        self.kernel = self.add_weight(
            "kernel",
            shape=[depth, self.output_units],
            # 以下は Optional
            dtype=tf.float32,
            initializer=tf.initializers.orthogonal(dtype=tf.float32),
        )
        # self.built = True と同義。最後に実行します。
        super().build(input_shape)

    def call(self, input):
        # 実行時の処理を書きます。
        return tf.matmul(input, self.kernel)

my_dense = MyDenseLayer(20)

input = tf.ones(shape=[2, 3])
output = my_dense(input)

# モデルの variable 一覧を取得できる！
print(my_dense.weights)  # -> [<tf.Variable 'my_dense_layer_4/kernel:0'...

罠

重みの宣言は必ず add_variable を使う

以前の tensorflow では重み宣言は　tf.get_variable や tf.Variable() で行っていましたが、 Layer の中では使わないほうが無難です。
（基礎から実践 TensorFlow 重み共有 で紹介した重み共有の方法は tensorflow 2.0 から廃止されるようです）

悪い例

class MyBadLayer(tf.keras.layers.Layer):
    def __init__(self, output_units):
        super().__init__()
        self.output_units = output_units

    def build(self, input_shape):
        depth = int(input_shape[-1])
        self.kernel = tf.get_variable(  # ここがダメ！ self.add_variable を使おう！
            "kernel", shape=[depth, self.output_units])
        super().build(input_shape)

    def call(self, input):
        # 実行時の処理を書きます。
        return tf.matmul(input, self.kernel)

l = MyBadLayer(10)
l(tf.ones(shape=[20, 30]))
l.weights  # -> []

上の例のように、レイヤーのインスタンスから重み一覧が参照できなくなってしまいます。

このような実装をしてしまうと、このレイヤーを使った Model で model.weights してそれを学習させてたつもりが一部の重みだけ学習されておらず学習がうまくいかない（けどエラーなく動く）という非常に気づきにくいバグを生むことになります。ひええ。
（tensorflow 2.0 からは optimizer へは重みの一覧を明示的に渡さないといけなくなるので上記のコードのように model.weights を使って学習対象の重みを指定することが多くなると予想されます。）

Layer は入れ子にしないほうがいい？

悪い例

class MyBadLayer2(tf.keras.layers.Layer):
    def build(self, input_shape):
        self.d = tf.keras.layers.Dense(20)  # これで layer.weights に追加されたと思っていたら・・
        super().build(input_shape)
        
    def call(self, input):
        return self.d(input)

l = MyBadLayer2()
l(tf.ones(shape=[20, 30]))
l.weights  # -> []  # 重み追加されてないじゃん

tensorflow 1.12 のコード を読んだ感じ、 add_variable しない限りは layer.weights に重みは追加されません。
したがって、 重みを持ちうる Layer を別の Layer の内部で使うのはおすすめしません。
後述する Model は入れ子にできるので混乱しないようにしましょう。

（違うよというのがあればご指摘お願いします。）

tf.keras.models.Model

自作した Layer や、 tensorflow が元々提供している Layer, また自作の別の Model などを組み合わせて Model を作りましょう。

雛形

• init: 部品となるレイヤーやモデルのインスタンスを作成します
• call: 上記をつなぎ合わせグラフを作ります

class Feedforward(tf.keras.models.Model):
    def __init__(self, hidden_dim_list, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dense_list = [tf.keras.layers.Dense(hidden_dim) for hidden_dim in hidden_dim_list]
         
    def call(self, inputs: tf.Tensor) -> tf.Tensor:
        tensor = inputs
        for dense in self.dense_list:
            tensor = dense(tensor)
        return tensor

m = Feedforward([20, 10, 2])
m(tf.ones(shape=[2, 30]))
m.weights  # -> [<tf.Variable 'feedforward_2/dense_20/...

モデルも Layer のように、自身の中に含まれる weight 一覧を取得できます。
しかも、自身が所持している（== self. でアクセスできる）レイヤーやモデルの中の weight もすべて合わせたものを出力してくれます。（次の節で説明します。）

階層的なモデル

Model は別のモデルを内部に持つこともできます。

class HierModel(tf.keras.models.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ff1 = Feedforward([10, 10])
        self.ff2 = Feedforward([5, 5])
        
    def call(self, input):
        return self.ff2(self.ff1(input))

m = HierModel()
m(tf.ones(shape=[10, 10]))
print('ff1:', len(m.ff1.weights))
print('ff2:', len(m.ff2.weights))
print('hier:', len(m.weights))
# -> ff1: 4
# -> ff2: 4
# -> hier: 8

このように、 hier は自身が内部に持っている ff1, ff2 両方の重みも返してくれます。
しかも、 m.ff1.weights などとすることでそのモデルの一部の重みのみを取ることもできます。

ネットワークを pre-train して、一部の重みのみをロードして transfer-learning するなども簡単にできそうですね。

罠

内部に持つモデルやレイヤーはインスタンスのメンバとして持とう

悪いモデル

class BadModel(tf.keras.models.Model):
    def __init__(self):
        super().__init__()
        # self.d = tf.keras.layers.Dense(20)  # こうしないとだめ
        
    def call(self, input):
        d = tf.keras.layers.Dense(20)  # レイヤーなどは self.** に入れないとだめ！
        return d(input)
        # or return tf.layers.dense(input, 20)  # つまり tf.layers.dense などは使えない！

m = BadModel()
out = m(tf.ones(shape=[2, 3]))
m.weights  # -> []

上の例でわかるように、重みを持ちうる layer や model はすべて、 self.** の形でアクセスできないといけません。
（コードを追ってみましたがおいきれなかったので正確でないかもしれません。）

以前の tensorflow 的な書き方だと tf.layers.dense （これも実は内部は tf.keras.layers.Dense で作られているけど） などをよく使うと思いますが、その流儀で作られたものはナウい Model では使えません。残念。正確には使えるのだけど、 weight 一覧を取ることができません。

レイヤーの集合をメンバで保つ場合、 List か Dict で

複数のレイヤーを保つ場合、素で持たせるか、 List, Dict を使いましょう。
Tuple を使うと tuple 内のレイヤーまでトラックしてくれないようです。

class SomeModel(tf.keras.models.Model):
    def __init__(self):
        super().__init__()
        # Tuple だと重みが追加されない！！！！！！
        self.my_layers = (tf.keras.layers.Dense(12), tf.keras.layers.Dense(11))

        # ↓のように List にしていればOK
        # self.my_layers = [tf.keras.layers.Dense(12), tf.keras.layers.Dense(11)]

    def call(self, input):
        for layer in self.my_layers:
            input = layer(input)
        return input

model = SomeModel()
model(tf.ones(shape=[1, 2, 3]))
model.weights  # -> []

最後に

いろいろなインターフェースが用意されているせいで罠がそこかしこにありますが、皆様の良い tf keras ライフをお祈りしております。
ちなみに Google 本家の Transformer とかのコードも Layer.build の中で tf.get_variable で変数宣言してたりしてどうなの・・？って気がします。