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を繰り返し作用させたいときなどに用います。 ([参考](https://www.tensorflow.org/api_docs/python/tf/keras/layers/TimeDistributed))

```python
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```っぽさを損なわずにすみます。([参考](https://www.tensorflow.org/tutorials/customization/custom_layers))

### 例.1: 標準の```tf.keras.layers```を組み合わせる
この場合、```.__init__()```で```tf.keras.layers```のinstanceを生成しておき、```.call()```にて実際のlayerの処理を定義します。

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

```python
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を使用する方法として、以下の様に、書けます。

```python
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```を使います。

```python
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
```
すると、以下のエラーが返ってきます。(一部抜粋)

```shell
/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()```というメソッドを用意しています。[コード](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/engine/base_layer.py/#L530)を見た限りbuildしようとしていますが、custom layerだとbuildに成功せずに、```NotImplementedError```がraiseしています。

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

```python
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```にわたしても正常に動いてくれます。