PythonとTensorFlowでカスタムレイヤーを作成・構成する方法

TensorFlowでは、tf.keras.Modelを継承したクラスを定義することで、独自のレイヤーやブロックを作成し、それらを組み合わせてニューラルネットワークを構築できます。本記事では、ResNetで使われる「残差ブロック（ResnetIdentityBlock）」を例に、カスタムレイヤーの作成方法を解説します。

前提知識

少なくとも1つの畳み込み層（Convolutional Layer）を含むニューラルネットワークは「畳み込みニューラルネットワーク（CNN）」と呼ばれ、画像認識などの学習モデルの構築に広く利用されています。

また、TensorFlow Hubは事前学習済みのTensorFlowモデルを集めたリポジトリです。tf.kerasと組み合わせることで、画像分類モデルなどをすぐに利用でき、転移学習によって独自の画像クラスに対応したモデルへファインチューニングすることも可能です。事前学習済みの分類器を使えば、追加のトレーニングなしに画像の予測を行うこともできます。

なお、本記事のコードはGoogle Colaboratoryでの実行を想定しています。Google Colabはブラウザ上でPythonコードを実行できる環境で、設定不要かつGPU（グラフィックス処理ユニット）へ無料でアクセスでき、Jupyter Notebookをベースに構築されています。

コード例：残差ブロックの定義

print("Composing layers")
class ResnetIdentityBlock(tf.keras.Model):
    def __init__(self, kernel_size, filters):
        super(ResnetIdentityBlock, self).__init__(name='')
        filters1, filters2, filters3 = filters
        self.conv2a = tf.keras.layers.Conv2D(filters1, (1, 1))
        self.bn2a = tf.keras.layers.BatchNormalization()
        self.conv2b = tf.keras.layers.Conv2D(filters2, kernel_size, padding='same')
        self.bn2b = tf.keras.layers.BatchNormalization()
        self.conv2c = tf.keras.layers.Conv2D(filters3, (1, 1))
        self.bn2c = tf.keras.layers.BatchNormalization()
    def call(self, input_tensor, training=False):
        x = self.conv2a(input_tensor)
        x = self.bn2a(x, training=training)
        x = tf.nn.relu(x)
        x = self.conv2b(x)
        x = self.bn2b(x, training=training)
        x = tf.nn.relu(x)
        x = self.conv2c(x)
        x = self.bn2c(x, training=training)
        x += input_tensor
        return tf.nn.relu(x)
print("The layer is called")
block = ResnetIdentityBlock(1, [1, 2, 3])
_ = block(tf.zeros([1, 2, 3, 3]))
block.layers
len(block.variables)
block.summary()

コード出典：https://www.tensorflow.org/tutorials/customization/custom_layers

実行結果

Composing layers
The layer is called
Model: "resnet_identity_block"
_________________________________________________________________
Layer (type)          Output Shape       Param #
=================================================================
conv2d (Conv2D)       multiple           4
_________________________________________________________________
batch_normalization (BatchNo multiple    4
_________________________________________________________________
conv2d_1 (Conv2D)     multiple           4
_________________________________________________________________
batch_normalization_1 (Batch multiple   8
_________________________________________________________________
conv2d_2 (Conv2D)     multiple           9
_________________________________________________________________
batch_normalization_2 (Batch multiple  12
=================================================================
Total params: 41
Trainable params: 29
Non-trainable params: 12

解説

• ResNet内のすべての残差ブロックは、畳み込み層（Convolution）、バッチ正規化（Batch Normalization）、そしてショートカット接続で構成されています。
• レイヤーは他のレイヤーの中に入れ子にすることもできます。
• Model.fit、Model.evaluate、Model.saveといったモデルのメソッドが必要な場合は、keras.Modelを継承します。
• keras.layers.Layerではなくkeras.Modelを使用することで、変数の追跡が容易になります。
• keras.Modelは内部のレイヤーを自動的に追跡するため、レイヤーの確認や検査が簡単に行えます。