Tensorflow2.0でLSTM+Attentionで分類モデルを作る

#はじめに
[TensorFlow2.0 Advent Calendar 2019](https://qiita.com/advent-calendar/2019/tensorflow2)の7日目のエントリで、tf.dataを使った[livedoorコーパス](https://www.rondhuit.com/download.html)の分かち書きとエンコーダ作成を行いました。

今エントリは前回の続きとして、`tf.keras.layers.LSTM`を使用してlivedoorコーパスの分類モデルを作成します。

#分類モデルについて
livedoorコーパスは全部で9つのジャンルに分かれていますが、今回は単純な分類モデルとしてテキストが`dokujo-tsushin`か否かの分類モデルを作成します。

#分かち書きデータセット作成
今回のデータセット作成は、大きく分けて以下の流れになります。

1. ラベル付きの分かち書きデータセットを用意
2. データセットからトークナイザ、エンコーダを作成
3. エンコーダを使用して分かち書きテキストをエンコード

### ラベル付きの分かち書きデータセットを用意

解凍済みlivedoorコーパスから`dokujo-tsushin`と`sports-watch`を残し、他をすべて削除します。

```
$ ls text/
dokujo-tsushin  sports-watch
```
[前回の記事](https://qiita.com/masahikoofjoyto/items/b444262405ad7371c78a#%E5%88%86%E3%81%8B%E3%81%A1%E6%9B%B8%E3%81%8D%E3%83%87%E3%83%BC%E3%82%BF%E3%82%BB%E3%83%83%E3%83%88%E3%82%92%E4%BD%9C%E6%88%90)を参考に、ラベル付きのデータセットを作成します。

`map`にテキストとラベルを渡すところが前回と異なっています。

```python
import os
import tensorflow as tf
import tensorflow_datasets as tfds
import fugashi
import numpy as np

tagger = fugashi.Tagger('-Owakati')

def parse(text):
    '''
    データセットのテキストを1行ごとに分かち書き
    '''
    return tagger.parse(text.numpy().decode('utf-8')).split("\n")[0]

@tf.function
def parse_text(text,label):
    parsed = tf.py_function(parse,[text],[tf.string])
    return parsed[0],label

text_datasets = []
text_dir = os.path.join(os.getcwd(),"text")
for d in tf.compat.v1.gfile.ListDirectory(text_dir):
    data_dir = os.path.join(text_dir,d)
    label = int(d == "dokujo-tsushin") # ディレクトリがdokujo-tsushinだったらTrue
    if os.path.isdir(data_dir):
        for file_name in tf.compat.v1.gfile.ListDirectory(data_dir):
            text_dataset = tf.data.TextLineDataset(os.path.join(data_dir,file_name)).map(lambda ex: parse_text(ex,label)) # 分かち書きコーパスとラベルのペアをデータセットに指定する
            text_datasets.append(text_dataset)

# テキスト単位の配列に変換する
parsed_dataset = text_datasets[0]
for text_dataset in text_datasets[1:]:
    parsed_dataset = parsed_dataset.concatenate(text_dataset)
```
分かち書き結果とラベルが返却されることを確認します。

```python
for txt,label in parsed_dataset.take(3):
    print('{} {}'.format(txt.numpy().decode('utf-8'), label))
```

```
http :// news . livedoor . com / article / detail / 6317287 / 1
2012 - 02 - 27 T 14 : 54 : 00 + 0900 1
話題 の 格言 『 ボーダー を 着る 女 は 、 95 ％ モテ ない ！ 』 は 本当 か ？ 1
```

### エンコーダ、トークナイザを用意
文字列をIDに変換するためのエンコーダとトークナイザを用意します。

トークナイザを用意します。

```python
tokenizer = tfds.features.text.Tokenizer()
```

トークナイザを使い、データセットからエンコーダを作成します。

```python
vocabulary_set = set()
for line,i in  parsed_dataset:
    some_tokens = tokenizer.tokenize(line.numpy())
    vocabulary_set.update(some_tokens)

encoder = tfds.features.text.TokenTextEncoder(vocabulary_set, tokenizer=tokenizer)
encoder.vocab_size
```

```
33784
```

# エンコーダを使用して分かち書きテキストをエンコード
データセットは`map`のチェーンが作れますので、`map`を使って分かち書きデータセットをエンコードします。

```python
def dataset_mapper(token,label):
    token = encoder.encode(token.numpy())
    label= np.array([label])
    return token,label

@tf.function
def tf_encode(token,label):
    return tf.py_function(dataset_mapper, [token,label], [tf.int64, tf.int64])

new_dataset = parsed_dataset.map(tf_encode)
```

バッチサイズを指定し、データをバッチサイズ単位に分割します。

```python
BATCH_SIZE = 128
new_dataset = new_dataset.padded_batch(BATCH_SIZE, padded_shapes=([-1],[-1]))
new_dataset = new_dataset.prefetch(tf.data.experimental.AUTOTUNE)
```

きちんと分割できているようです。

```python
text_batch, label_batch = next(iter(new_dataset))
print(text_batch.shape)
print(label_batch.shape)
```

```
(128, 155)
(128, 1)
```

# LSTMモデルの作成
データセットの準備ができたので、モデル作成に入ります。

Tensorflow2.0はkerasと統合されているので、kerasのわかりやすいインターフェイスで書けるのが素晴らしいです。

```python
X = tf.keras.Input(shape=(None,), batch_size=BATCH_SIZE)
embedded = tf.keras.layers.Embedding(encoder.vocab_size, 128)(X)
lstm = tf.keras.layers.Bidirectional(
      tf.keras.layers.LSTM(128, dropout=0.4, recurrent_dropout=0.4)
  )(embedded)
fully_connected = tf.keras.layers.Dense(units=256, activation='relu')(lstm)
Y = tf.keras.layers.Dense(1, activation='sigmoid',name='final_layer')(fully_connected)

model = tf.keras.Model(inputs=X, outputs=Y)

model.compile(loss='binary_crossentropy',
              optimizer=tf.keras.optimizers.Adam(1e-4),
              metrics=['accuracy'])
```

モデルのレイヤを確認します。

```python
model.summary()
```

```
Model: "model_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(128, None)]             0         
_________________________________________________________________
embedding_1 (Embedding)      (128, None, 128)          4324352   
_________________________________________________________________
bidirectional_1 (Bidirection (128, 256)                263168    
_________________________________________________________________
dense_1 (Dense)              (128, 256)                65792     
_________________________________________________________________
final_layer (Dense)          (128, 1)                  257       
=================================================================
Total params: 4,653,569
Trainable params: 4,653,569
Non-trainable params: 0
_________________________________________________________________
```

モデルが完成したので、早速学習を開始します。

```python
model.fit(new_dataset,epochs=5)
```

```
Epoch 1/5
298/298 [==============================] - 140s 471ms/step - loss: 0.5049 - accuracy: 0.8531
Epoch 2/5
298/298 [==============================] - 137s 459ms/step - loss: 0.7768 - accuracy: 0.5536
Epoch 3/5
298/298 [==============================] - 245s 824ms/step - loss: 0.6405 - accuracy: 0.6531
Epoch 4/5
298/298 [==============================] - 327s 1s/step - loss: 0.5248 - accuracy: 0.8055
Epoch 5/5
298/298 [==============================] - 417s 1s/step - loss: 0.4686 - accuracy: 0.8019
```
エポック数が少なすぎるため、あまり精度の出ませんでしたが気にせず次に進みます。

# Attention
[Tensorflow 2.0チュートリアル](https://www.tensorflow.org/tutorials)はサンプルと知見の宝庫でとても素晴らしく、チュートリアルの[Attention](https://www.tensorflow.org/tutorials/text/nmt_with_attention#write_the_encoder_and_decoder_model)実装を参考にレイヤを作成します。
チュートリアルにはAttentionはもちろん、他にも様々なタスクの実装サンプルが記述されており、有志の方々が翻訳された日本語版も大変わかりやすいのでまだ読んでいない方は是非ご一読ください。


```python
class Attention(tf.keras.Model):
    def __init__(self, units):
        super(Attention, self).__init__()
        self.W1 = tf.keras.layers.Dense(units)
        self.W2 = tf.keras.layers.Dense(units)
        self.V = tf.keras.layers.Dense(1)

    def call(self, features, hidden):
        hidden_with_time_axis = tf.expand_dims(hidden, 1)
        score = tf.nn.tanh(self.W1(features) + self.W2(hidden_with_time_axis))
        attention_weights = tf.nn.softmax(self.V(score), axis=1)
        context_vector = attention_weights * features
        context_vector = tf.reduce_sum(context_vector, axis=1)

        return context_vector, attention_weights
```

先ほどのLSTMレイヤを改良し、Attentionレイヤを挟みます。LSTMは双方向LSTMとしたので、今回は順・逆両方の重みを結合したものをAttentionレイヤのinputとしています。

```python
X = tf.keras.Input(shape=(None,), batch_size=BATCH_SIZE)
embedded = tf.keras.layers.Embedding(encoder.vocab_size, 128)(X)
lstm, forward_h, forward_c, backward_h, backward_c = tf.keras.layers.Bidirectional(
      tf.keras.layers.LSTM(128,return_sequences=True,return_state=True, dropout=0.4, recurrent_dropout=0.4)
  )(embedded)
state_h = tf.keras.layers.Concatenate()([forward_h, backward_h]) # 重みを結合
context,attention_weights = Attention(128)(lstm,state_h) # ここにAttentionレイヤを挟む
fully_connected = tf.keras.layers.Dense(units=256, activation='relu')(context)
Y = tf.keras.layers.Dense(1, activation='sigmoid',name='final_layer')(fully_connected)

model = tf.keras.Model(inputs=X, outputs=Y)

model.compile(loss='binary_crossentropy',
              optimizer=tf.keras.optimizers.Adam(1e-4),
              metrics=['accuracy'])
```

グラフを表示し、Attentionレイヤが間に入っていることを確認します。

```pyton
model.summary()
```

```
Model: "model"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(128, None)]        0                                            
__________________________________________________________________________________________________
embedding (Embedding)           (128, None, 128)     4324352     input_1[0][0]                    
__________________________________________________________________________________________________
bidirectional (Bidirectional)   [(128, None, 256), ( 263168      embedding[0][0]                  
__________________________________________________________________________________________________
concatenate (Concatenate)       (None, 256)          0           bidirectional[0][1]              
                                                                 bidirectional[0][3]              
__________________________________________________________________________________________________
attention (Attention)           ((128, 256), (128, N 65921       bidirectional[0][0]              
__________________________________________________________________________________________________
dense_3 (Dense)                 (128, 256)           65792       attention[0][0]                  
__________________________________________________________________________________________________
final_layer (Dense)             (128, 1)             257         dense_3[0][0]                    
==================================================================================================
Total params: 4,719,490
Trainable params: 4,719,490
Non-trainable params: 0
__________________________________________________________________________________________________
```

Attentionレイヤ付きのモデルが完成したので、学習させます。

```python
model.fit(new_dataset,epochs=5)
```

精度は悪くなってますが、お財布と時間の関係でこれ以上学習できないため、ここで打ち切りとします。。。

```
Epoch 1/5
298/298 [==============================] - 123s 413ms/step - loss: 0.4978 - accuracy: 0.8718
Epoch 2/5
298/298 [==============================] - 121s 405ms/step - loss: 0.7790 - accuracy: 0.4971
Epoch 3/5
298/298 [==============================] - 121s 406ms/step - loss: 0.6752 - accuracy: 0.6168
Epoch 4/5
298/298 [==============================] - 121s 405ms/step - loss: 0.6523 - accuracy: 0.6869
Epoch 5/5
298/298 [==============================] - 121s 406ms/step - loss: 0.6395 - accuracy: 0.6895
```

# まとめ
`tf.data`と`tf.keras`を使うと簡単にLSTMレイヤのモデルが作れます。

また今回のように独自モデルを加えることも容易にできますので、ぜひ試してみてください。