🌭

TensofFlowの使い方練習４：サブクラス化

2023/04/06に公開

Python

TensorFlow

tech

はじめに

この記事は以下の記事の続きである。

この記事では tensorflow.keras.layers.Layer や tensorflow.keras.Model を継承して新たなレイヤーやモデルを作成する方法を練習する。

参考にしたドキュメントは以下。

参考文献

独自レイヤーの定義

独自のレイヤーを定義するには tensorflow.keras.layers.Layer を継承し、__init__() メソッドと call() メソッドをオーバーライドする（__call__()ではないことに注意）。

レイヤーが内部に状態を持つ場合は自身のメンバとして追加する。

import tensorflow as tf
from tensorflow import keras

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__init__()
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(
            initial_value=w_init(shape=(input_dim, units), dtype="float32"),
            trainable=True,
        )
        b_init = tf.zeros_initializer()
        self.b = tf.Variable(
            initial_value=b_init(shape=(units,), dtype="float32"), trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

このように定義されたレイヤーは Keras に実装されているレイヤーと同じようにインスタンス化して呼び出すことができる。

import numpy as np

# 擬似データ
x = np.random.normal(0, 1, (32, 32))

layer = Linear()

layer(x)

メンバとして追加された状態は Layer クラスの機能により自動で追跡される。追跡対象となっている状態は weights メンバで取得できる。

print(layer.weights)

output

[<tf.Variable 'Variable:0' shape=(32, 32) dtype=float32, numpy=
 array([[ 0.10401148,  0.00977317, -0.04388207, ...,  0.115283  ,
         -0.04799509, -0.03706598],
        [ 0.02498189,  0.05101189, -0.04151009, ...,  0.02203462,
         -0.02765396,  0.08929577],
        [ 0.00127457, -0.05099274,  0.08889019, ..., -0.02877973,
         -0.08359879,  0.04306518],
        ...,
        [-0.01639055,  0.0687128 ,  0.01887065, ...,  0.0531274 ,
          0.01022338, -0.00353025],
        [-0.02771079,  0.07625233,  0.00402813, ..., -0.01826211,
         -0.01207191,  0.03357437],
        [ 0.06055077, -0.03304816,  0.1236055 , ...,  0.09350461,
          0.03379982, -0.02095563]], dtype=float32)>,
 <tf.Variable 'Variable:0' shape=(32,) dtype=float32, numpy=
 array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       dtype=float32)>]

ここで気になるのは

文字列などをメンバとして追加しても weights として追跡されてしまうのか？
計算に直接関わらない、たとえば self.a = ... を追加しても追跡されるのか？

である。実際にやってみると前者は No, 後者は Yes である。挙動としては

テンソル（tf.Variable）がメンバとして追加されればそれが計算に使われるかどうかに関わらず追跡対象になる。
それ以外は追跡対象にならない。

である。また、テンソルがメンバとして追加される場合はプライベートメンバ（名前の先頭が _ または __ で始まるメンバ）であるかどうかに関わらず追跡対象になる。

この挙動に関わりそうなのは Layer クラスの特殊メソッドの __setattr__() とプロパティであるはずの weights である。実際にソースコードを見に行く。

`setattr()` メソッド

セットされたメンバが tf.Variable のインスタンスであれば __setattr__()　メソッド内の _track_variables() で追跡対象に追加されているようである。__setattr__() を見るとその他様々な対象を追跡対象にしているようである。

重みの追加に限定するのであれば add_weights メソッドを使って直接重みを追加することもできる。見た目は __setattr__() を使ったほうが綺麗かもしれないが、add_weights を使うと重みの追加に際していろいろな処理が行われているだろうことを明示できるのでこの辺りは好みだろう。

`weights` プロパティ

__setattr__() で追跡対象と判断された tf.Variable は trainable_weights と non_trainable_weights に分けて管理される。どちらになるかは追加される tf.Variable が trainable=True/False のどちらで生成されているかによる。

`build` メソッド

呼び出し時までレイヤーが持つ状態の shape が分からない、または呼び出し時に状態の shape を決めたほうが都合がいい場合には build メソッドを定義しておく。

build メソッドは __call__() メソッドが最初に呼び出されたときに呼び出される。つまり与えられた入力の shape に応じてレイヤーを初期化することができる。

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

`get_config` メソッド

Functional API に統合するための Serialization を有効にするには get_config() メソッドも定義する。

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

    def get_config(self):
        return {"units": self.units}


# Now you can recreate the layer from its config:
layer = Linear(64)
config = layer.get_config()
print(config)
new_layer = Linear.from_config(config)

レイヤーのネスト

レイヤーはネストできる。

class MLPBlock(keras.layers.Layer):
    def __init__(self):
        super(MLPBlock, self).__init__()
        self.linear_1 = Linear(32)
        self.linear_2 = Linear(32)
        self.linear_3 = Linear(1)

    def call(self, inputs):
        x = self.linear_1(inputs)
        x = tf.nn.relu(x)
        x = self.linear_2(x)
        x = tf.nn.relu(x)
        return self.linear_3(x)

mlp = MLPBlock()
y = mlp(tf.ones(shape=(3, 64)))  # The first call to the `mlp` will create the weights
print("weights:", len(mlp.weights))
print("trainable weights:", len(mlp.trainable_weights))

学習時と推論時で挙動を切り替える

call() メソッドの引数の training は学習中であれば True になる。

class CustomDropout(keras.layers.Layer):
    def __init__(self, rate, **kwargs):
        super(CustomDropout, self).__init__(**kwargs)
        self.rate = rate

    def call(self, inputs, training=None):
        if training:
            return tf.nn.dropout(inputs, rate=self.rate)
        return inputs

損失関数とメトリックを追加する

使うかどうかわからないのでここには書かないが add_loss() メソッドで損失関数を追加することでレイヤー単位での損失関数を追加することができる。たとえば重みのノルムとかをモデル全体ではなくレイヤー単位で追加しておくことで、最適化のときにモデル全体の損失関数に足されて最適化の対象となる。

同様に add_metric() は Accuracy のような評価指標を追加しておくことができる。

add_loss() メソッドや add_metric() メソッドは call() メソッドの中に記述し、実行される度に損失や評価指標を更新するように用いられる。

独自モデルの定義

レイヤーを積み重ねたのがモデル。モデルがガワでレイヤーが中身。人間でいうと人間そのものがモデル。その内臓がレイヤー。

モデルはレイヤーと比べて

fit() メソッドがあり訓練することができる
save() メソッドで自身を保存することができる

という点が異なる。それ以外の違いはあまりない……のか？サンプルプログラムは以下のようなものが載せられている。

class ResNet(tf.keras.Model):

    def __init__(self, num_classes=1000):
        super(ResNet, self).__init__()
        self.block_1 = ResNetBlock()
        self.block_2 = ResNetBlock()
        self.global_pool = layers.GlobalAveragePooling2D()
        self.classifier = Dense(num_classes)

    def call(self, inputs):
        x = self.block_1(inputs)
        x = self.block_2(x)
        x = self.global_pool(x)
        return self.classifier(x)

resnet = ResNet()
dataset = ...
resnet.fit(dataset, epochs=10)
resnet.save(filepath)

Layer の拡張みたいなものなのだとしたら Model は Layer を継承していて何か付け加えたもののはず。ほんとにそうなのかコードを確かめに行こう。

継承してる。思ったよりもたくさんの機能が追加されていて、save, fit などの他にも compile などのメソッドが追加されており build などが実行されるタイミングも統制しているように見える。

訓練に関するメソッドなどもほとんどモデルのほうに書かれていて、機械学習モデルとして機能させるには tf.Model になっていないといけないようである。

まとめ

ここまで見てきて tf.Variable → tf.keras.layers.Layer → tf.Model の順に大きくなって機能が付加されていくことがわかった。

今シリーズ、チュートリアルをなぞっているだけの部分が多く価値のない記事だし、タイトルに「練習」とかつけてしまったせいか、コミュニティガイドライン的に不適切とみなされているようで露骨にリアクションが少ない。

基本的な部分は押さえた気がするのでいったん TensorFlow 使い方練習シリーズはこれで終わりにして、次回からはいろんなモデルを作ったり解剖したりして遊ぶことにする。