tf.kerasのcustom layerに対してtf.keras.layers.TimeDistributedを使う時の注意

本記事の目的

tf.kerasで定義したcustom layerに対してtf.keras.layers.TimeDistributedを用いるとエラーを吐いたので、その内容と解決策を共有します。

version

• Python: 3.6.9
• Tensorflow: 2.1.0

目次

• tf.keras.layers.TimeDistributedとは
• custom layerとは
• custom layerに対してtf.keras.layers.TimeDistributedを作用させたときのエラー内容紹介
• 上記の解決法

tf.keras.layers.TimeDistributedとは

時間方向に1つのlayerを繰り返し作用させたいときなどに用います。 (参考)

from tensorflow.keras import Input, Model
from tensorflow.keras.layers import TimeDistributed, Dense

temporal_dim = 4
emb_dim = 16
inputs = Input((temporal_dim, 8))
outputs = TimeDistributed(Dense(emb_dim))(inputs)
model = Model(inputs, outputs)
model.summary()
"""
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 4, 8)]            0         
_________________________________________________________________
time_distributed (TimeDistri (None, 4, 16)             144       
=================================================================
Total params: 144
Trainable params: 144
Non-trainable params: 0
_________________________________________________________________
"""

注) tf.keras.layers.Denseは最後の次元にしか作用しないので、上記結果はtf.keras.layers.TimeDistributedの有無に依らないです。ありがたみを感じられるのは、tf.keras.layers.Embeddingのようなinput_shapeに制限があるlayerに対して使った時です。

tf.keras.layers.Embeddingはbatch_sizeこみで2D tensorにしか使えないので、もしも、3D tensorを埋め込みたい場合は、tf.keras.layers.TimeDistributedを使って足りない余計な1次元方向に繰り返し作用させる必要があります。(recommendのモデルで使います)

custom layerとは

tf.kerasでは、tf.keras.layers.Layerを継承してcustom layerを定義できます。
tf.kerasにはたくさんの便利なlayerが実装されており、それらを組み合わせてtf.keras.Modelを作り上げていきます。標準のlayerだけでも非常に多くのモデルが実装できます。

より柔軟な処理を行いたい場合には、custom layerとして実装することでtf.kerasっぽさを損なわずにすみます。(参考)

例.1: 標準のtf.keras.layersを組み合わせる

この場合、.__init__()でtf.keras.layersのinstanceを生成しておき、.call()にて実際のlayerの処理を定義します。

（以下のlayerは、tf.keras.layers.Denseと同じものです）

import tensorflow as tf
from tensorflow.keras.layers import Dense

class Linear(tf.keras.layers.Layer):
    def __init__(self, emb_dim, *args, **kwargs):
        super(Linear_error, self).__init__(*args, **kwargs)
        self.dense = Dense(emb_dim)

    def call(self, x):
        return self.dense(x)

例.2: 学習可能な変数を定義する

より低レベルAPIを使用する方法として、以下の様に、書けます。

class Linear(tf.keras.layers.Layer):
    def __init__(self, emb_dim):
        super(Linear, self).__init__()
        self.emb_dim = emb_dim

    def build(self, input_shape):
        self.kernel = self.add_weight(
            "kernel", shape=[int(input_shape[-1]), self.emb_dim]
        )

    def call(self, x):
        return tf.matmul(x, self.kernel)

Error: custom layerに対するtf.keras.layers.TimeDistributedが動かない

以下の様にcustom layerに対してtf.keras.layers.TimeDistributedを使います。

from tensorflow.keras import Input
from tensorflow.keras.layers import TimeDistributed

temporal_dim = 4
emb_dim = 16
inputs = Input((temporal_dim, 8))

outputs = TimeDistributed(Linear(emb_dim))(inputs)
assert outputs.shape.rank == 3
assert outputs.shape[1] == temporal_dim
assert outputs.shape[2] == emb_dim

すると、以下のエラーが返ってきます。(一部抜粋)

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/base_layer.py:773: in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/layers/wrappers.py:270: in call
    output_shape = self.compute_output_shape(input_shape).as_list()
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/layers/wrappers.py:212: in compute_output_shape
    child_output_shape = self.layer.compute_output_shape(child_input_shape)

>     raise NotImplementedError
E     NotImplementedError

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/base_layer.py:564: NotImplementedError

これは、custom layerの.compute_output_shape()の中身がないので自分で実装して下さいというメッセージです。

tf.keras.layersは.build()か、.call()が呼ばれるまでinput_shapeとoutput_shapeが定まらないです。しかし、tf.keras.layers.TimeDistributedは引数に入れたlayerのouput_shapeをinstance生成時(つまり、.build()や.call()の前)に知ろうとします。tf.keras.layersはそんな時のために.compute_output_shape()というメソッドを用意しています。コードを見た限りbuildしようとしていますが、custom layerだとbuildに成功せずに、NotImplementedErrorがraiseしています。

解決策

custom layerの定義時に.compute_output_shape()をoverrideして陽にoutput_shapeを定義します。

class Linear(tf.keras.layers.Layer):
    def __init__(self, emb_dim, *args, **kwargs):
        super(Linear, self).__init__(*args, **kwargs)
        self.emb_dim = emb_dim
        self.dense = Dense(self.emb_dim)

    def call(self, x):
        return self.dense(x)

    def compute_output_shape(self, input_shape):
        output_shape = input_shape[0:-1] + [self.emb_dim]
        return output_shape

こちらのcustom layerであれば、tf.keras.layers.TimeDistributedにわたしても正常に動いてくれます。