論文の勉強15 Transformer(Encoderのみ)

Transformerについて構造の説明と実装のメモ書きです。
ただし、論文すべてを見るわけでなく構造のところを中心に見ていきます。

transformerは画像認識分野でも使われるようになっています。
ここでは分類モデルを意識して**Encoder部分**のみの実装となります。

勉強のメモ書き程度でありあまり正確に実装されていませんので、ご了承ください。
自分の実力不足で読み解けなくなってきています。難しいです。

以下の論文について実装を行っていきます。

タイトル：Attentino Is All You Need

https://arxiv.org/abs/1706.03762

以下のサイトや本を参考にしています。

https://qiita.com/halhorn/items/c91497522be27bde17ce  

https://www.amazon.co.jp/%E4%BD%9C%E3%81%A3%E3%81%A6%E3%82%8F%E3%81%8B%E3%82%8B-%E8%87%AA%E7%84%B6%E8%A8%80%E8%AA%9E%E5%87%A6%E7%90%86AI%E3%80%9CBERT%E3%83%BBGPT2%E3%83%BBNLP%E3%83%97%E3%83%AD%E3%82%B0%E3%83%A9%E3%83%9F%E3%83%B3%E3%82%B0%E5%85%A5%E9%96%80-%E5%9D%82%E6%9C%AC-%E4%BF%8A%E4%B9%8B/dp/4863543700/ref=pd_sim_sccl_6/357-1227394-8103025?pd_rd_w=qNnkm&pf_rd_p=ba248832-b816-4425-8851-893f63228155&pf_rd_r=ASB2H0TH27YWHE6DN9M6&pd_rd_r=77502e62-7bc5-4152-9dd1-a28864140f0e&pd_rd_wg=HT1rn&pd_rd_i=4863543700&psc=1

### Attention

Encoder-Decoderモデルとなっています。
Encoderでは入力$(x_1,\cdots,x_n)$を中間表現$(z_1,\cdots,z_n)$に変換します。
Decoderでは、それを受け取って出力$(y_1,\cdots,y_n)$を生成します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2358656/92c8b6ef-5d13-3284-3044-1e5c5adec85a.png)


#### Encoder
$N=6$個の独立したレイヤで構成されます。
それぞれのレイヤは2つのサブレイヤを持ち、1つはmulti-head self-attentionであり、2つ目は全結合のfeed-forward networkとなります。
それぞれ残差接続(residual connection)を持ち、加算後はLayer Normalizationを適用します。
すべてのサブレイヤの出力の次元は$d_{model}=512$となります。

#### Decoder
$N=6$個の独立したレイヤで構成されます。
Encoder同様の2つのサブレイヤに加えて、3つ目のサブレイヤが挿入されています。
これはEncoderの出力を受け取るmulti-head attentionです。
Decoderにおいては、未来の情報を入れないようにmask処理をします。
この処理は、図でいうとmasked multi-head attentionで行われます。

#### Attention
attentionはqueryとkey-valueのセットのマッピングとして表されます。  
query,key,value,そして出力はすべてベクトルとなります。
出力はvakueのwieghted sum(加重和)であり、重みはqueryとkeyから計算されます。

#### Scaled Dot-Product Attention

下の図で表されるScaled Dot-Product Attentionを考えます。
queryとkeyの次元は$d_k$、valueの次元は$d_v$です。
queryとkeyのドット積を$\sqrt{d_k}$で割ったものをsoftmax関数に入れることで重みを計算します。
ベクトルであったquery,key,valueをそれぞれまとめて行列$Q,K,V$と表します。
そしてScaled Dot-Product Attentionの出力を
$$
Attention(Q,K,V)=softmax\bigl(\frac{QK^T}{\sqrt{d_k}} \bigr)V
$$
と計算します。

#### Multi-Head Attention
query,key,valueがすべて$d_{model}$次元で処理するより、線形射影によりそれぞれ$h_k,h_k,h_v$次元にした方が効果的であることが分かりました。これらは$d_{model}$の$h$分の1となっています。
それぞれAttention層で並行して処理をして$d_v$次元の出力を得ます。これらを結合して最終的な出力とします。

$$
MultiHead(Q,K,V)=Concat(head_1,\cdots,head_h)W^O\\
where\ head_i=Attention(QW_i^Q,KW_i^K,VW_i^V)
$$
ここで、$W_i^Q\in R^{d_{model}×d_k},W_i^K\in R^{d_{model}×d_k},W_i^V\in R^{d_{model}×d_v}$、そして$W^O=R^{hd_v×d_{model}}$です。
今回は$h=8$で、さらに$d_k=d_v=d_{model}/h=64$とします。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2358656/f88d06d7-1d7b-4005-30c5-4de03b010cda.png)


#### Position-wise Feed-Foward Networks
もう1つのサブレイヤであるFeed-Foward Networksを説明します。
このレイヤではすべての位置に対して独立に全結合層が適用されます。
つまり1×1のConvolution層が使用されます。
ここでは2つの全結合層を使用して、1つ目のConv層で$d_{ff}=2048$、2つ目のConv層で$d_{model}=512$へと次元を変換します。
また、1つ目のConv層ではReLU関数を使用します。
$$
FFN(x)=max(0,xW_1+b)W_2+b_2
$$

#### Embedding and Softmax
inputとoutputに対して、$d_{model}$次元へ変換する同じ学習済みの重みを使用します。
decoderの出力に対してはSoftmax関数により次の単語の出現確率を計算します。ここでもinput,outputと同じ重みを使用します。
また、embeddingレイヤでは重みに$\sqrt{d_{model}}$を掛けたものを使用します。

#### Positional Encoding
RNN・CNNがないため、単語の位置情報を与える必要があります。
input.outputのEmbedding層の後にpositional encodingを追加します。
これは$d_{model}$の次元を持ち、embeddingされたinputやoutputと足し合わせます。
様々なものがありますがここでは,
$$
PE_{(pos,2i)}=\sin{(pos/10000^{2i/d_{model}})}\\
PE_{(pos,2i+1)}=\cos{(pos/10000^{2i/d_{model}})}
$$
を使用します。
$pos$は位置を表し、$i$は次元を表します。


#### 学習
最適化手法としてAdamを使用し、$\beta_1=0.9,\beta_2=0.98,\epsilon=10^{-9}$です。
学習率は
$$
lrate=d_{model}^{-0.5}\min(step_num^{-0.5},step_num・warmup_steps^{-1.5})
$$
に従うものとします。
warmup_stepsでは学習率は増加し、その後減少していきます。
今回はwaramup_steps=4000としました。

#### 正則化(dropout・label smoothing)
各サブレイヤの後にdropout層を入れました。また、encoder・decoderそれぞれでembedding・positionnal encodingのあとにもDropoutを入れています。
確率は$P_{drop}=0.1$とします。

label smoothingも採用しています。
$\epsilon_{ls}=0.1$とします。


### keras

```python
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Layer, Input, Dense, Conv1D, Activation, Dropout, LayerNormalization, Reshape, Embedding, MultiHeadAttention
from tensorflow.keras import activations
from tensorflow.keras.preprocessing.sequence import pad_sequences
import numpy as np
import math
```

MultiHeadAttentionを実装します。
可視化のためにAttention weightの出力を選択できるようにします。

```python
class MultiHeadAttention(Layer):
    '''
    Multi-Head Attentionレイヤ
    
    hidden_dim : Embeddingされた単語ベクトルの長さ
    heads_num : マルチヘッドAttentionのヘッド数
       ※hidden_numはheads_numで割り切れえる値とすること
    drop_rate : 出力のDropout率

    model = MultiheadAttention(
        hidden_dim = 512,
        head_num = 8,
        drop_rate = 0.5
    )
    '''
    def __init__(self, hidden_dim, heads_num, drop_rate=0.5):
        super(MultiHeadAttention, self).__init__()
        # 入力の線形変換
        # 重み行列は[hidden_dim, hidden_dim]
        self.query = Conv1D(hidden_dim, kernel_size=1)
        self.key   = Conv1D(hidden_dim, kernel_size=1)
        self.value = Conv1D(hidden_dim, kernel_size=1)
        
        # 出力の線形変換
        self.projection = Conv1D(hidden_dim, kernel_size=1)
        
        # 出力のDropout
        self.drop = Dropout(drop_rate)
        
        self.nf = hidden_dim
        self.nh = heads_num
    
    def atten(self, query, key, value, attention_mask, training):
        """
        Attention
        
        query, key, value : 入力
        attention_mask : attention weight に適用される mask
        """
        # 各値を取得
        shape = query.shape.as_list()
        batch_size = -1 if shape[0] is None else shape[0]
        token_num = shape[2] # トークン列数
        hidden_dim = shape[1]*shape[3] # 入力チャンネル数
        
        # ここで q と k の内積を取ることで、query と key の単語間の関連度のようなものを計算します。
        # tf.matmulで最後の2成分について積を計算(それ以外は形がそろっている必要あり)
        # transpose_bで転置
        # [token_num, hidden_dim/head_num] @ [hidden_dim/head_num, token_num] = [token_num, token_num]
        scores = tf.matmul(query, key, transpose_b=True)
        
        # scoreをhidden_dimの平方根割る
        scores = tf.multiply(scores, tf.math.rsqrt(tf.cast(hidden_dim, tf.float32)))
        
        # Attention Maskがあればscoreに加算
        # attention_mask: [batch_size, token_num, token_num] 
        # マスク(参照しない部分)の場所に1、使用する部分は0とする
        # 0の部分を -無限大にする(softmax(-無限大)=0となる)
        # 1. PADを無視
        # 2. DecoderのSelf-Attentionで未来の情報を参照できないようにする
        if attention_mask is not None:
            scores += attention_mask * -1e9
        # softmax を取ることで正規化します
        # input(query) の各単語に対して memory(key) の各単語のどこから情報を引いてくるかの重み
        atten_weight = tf.nn.softmax(scores)
        
        # 重みに従って value から情報を引いてきます
        # [token_num, token_num] @ [token_num, hidden_dim/head_num] = [token_num, hidden_dim/head_num]
        # input(query) の単語ごとに memory(value)の各単語 に attention_weight を掛け合わせて足し合わせた ベクトル(分散表現の重み付き和)を計算
        context = tf.matmul(atten_weight, value)
        
        # 各ヘッドの結合(reshape)
        # 入力と同じ形に変換する
        context = tf.transpose(context, [0, 2, 1, 3])
        context = tf.reshape(context, (batch_size, token_num, hidden_dim))
        
        # 線形変換
        context = self.projection(context, training=training)
        
        return self.drop(context, training=training), atten_weight

    def _split(self, x):
        """
        query, key, valueを分割する
        
        入力 shape: [batch_size, length, hidden_dim] の時
        出力 shape: [batch_size, head_num, length, hidden_dim//head_num]
        """
        # 各値を取得
        hidden_dim = self.nf
        heads_num = self.nh
        shape = x.shape.as_list()
        batch_size = -1 if shape[0] is None else shape[0]
        token_num = shape[1] # トークン列数
        
        # [batch_size, token_num, hidden_dim] -> [batch_size, token_num, head_num, hidden_dim/head_num]
        # splitだが実際は次元を拡張する処理
        x = tf.reshape(x, (batch_size, token_num, heads_num, int(hidden_dim/heads_num)))
        
        # [batch_size, token_num, head_num, hidden_dim/head_num] -> [batch_size, head_num, token_num, hidden_dim/head_num]
        x = tf.transpose(x, [0, 2, 1, 3])
        return x
    
    def call(self, x, training, memory=None, attention_mask=None, return_attention_scores=False):
        """
        モデルの実行
        
        input : 入力(query) [batch_size, token_num, hidden_dim]
        memory : 入力(key, value) [batch_size, token_num, hidden_dim]
        attention_mask : attention weight に適用される mask
            [batch_size, 1, q_length, k_length] 
            pad 等無視する部分が 1 となるようなもの(Decoderで使用)
        return_attention_scores : attention weightを出力するか
        """
        # memoryが入力されない場合、memory=input(Self Attention)とする
        if memory is None:
            memory = x
        
        # input -> query
        # memory -> key, value
        # [batch_size, token_num, hidden_dim] @ [hidden_dim, hidden_dim] -> [batch_size, token_num, hidden_dim] 
        query = self.query(x)
        key = self.key(memory)
        value = self.value(memory)
        
        # ヘッド数に分割する
        # 実際はreshapeで次数を1つ増やす
        # [batch_size, token_num, hidden_dim] -> [batch_size, head_num, token_num, hidden_dim/head_num]
        query = self._split(query)
        key = self._split(key)
        value = self._split(value)
        
        # attention
        # 入力と同じ形の出力
        # context: [batch_size, token_num, hidden_dim]
        # score_weightsはEncoderではNoneとする
        context, attn_weights = self.atten(query, key, value, attention_mask, training)
        if not return_attention_scores:
            return context
        else:
            return context, attn_weights
```

FeedForwardNetworkの実装をします。
2つのDense Layerからなります。

```python
class FeedForwardNetwork(Layer):
    '''
    Position-wise Feedforward Neural Network
    transformer blockで使用される全結合層
    '''
    def __init__(self, hidden_dim, drop_rate):
        super().__init__()
        # 2層構造
        # 1層目：チャンネル数を増加させる
        self.filter_dense_layer = Dense(hidden_dim * 4, use_bias=True, activation='relu')
        
        # 2層目：元のチャンネル数に戻す
        self.output_dense_layer = Dense(hidden_dim, use_bias=True)
        self.drop = Dropout(drop_rate)

    def call(self, x, training):
        '''
        入力と出力で形が変わらない
        [batch_size, token_num, hidden_dim]
        '''
        
        # [batch_size, token_num, hidden_dim] -> [batch_size, token_num, 4*hidden_dim]
        x = self.filter_dense_layer(x)
        x = self.drop(x, training=training)
        
        # [batch_size, token_num, 4*hidden_dim] -> [batch_size, token_num, hidden_dim]
        return self.output_dense_layer(x)
```

残差接続のためのResidualNormalizationWrapperを実装します。

```python
class ResidualNormalizationWrapper(Layer):
    '''
    残差接続
    output: input + SubLayer(input)
    '''
    def __init__(self, layer, drop_rate):
        super().__init__()
        self.layer = layer # SubLayer : ここではAttentionかFFN
        self.layer_normalization = LayerNormalization()
        self.drop = Dropout(drop_rate)

    def call(self, x, training, memory=None, attention_mask=None, return_attention_scores=None):
        """
        AttentionもFFNも入力と出力で形が変わらない
        [batch_size, token_num, hidden_dim]
        """
        
        params = {}
        if memory is not None:
            params['memory'] = memory
        if attention_mask is not None:
            params['attention_mask'] = attention_mask
        if return_attention_scores:
            params['return_attention_scores'] = return_attention_scores
        
        out = self.layer_normalization(x)
        if return_attention_scores:
            out, attn_weights = self.layer(out, training, **params)
            out = self.drop(out, training=training)
            return x + out, attn_weights
        else:
            out = self.layer(out, training, **params)
            out = self.drop(out, training=training)
            return x + out
```

PositionalEncodingを行うレイヤの実装をします。
いろいろ実装方法はあるようです。

```python
class AddPositionalEncoding(Layer):
    '''
    入力テンソルに対し、位置の情報を付与して返すレイヤー
    see: https://arxiv.org/pdf/1706.03762.pdf

    PE_{pos, 2i}   = sin(pos / 10000^{2i / d_model})
    PE_{pos, 2i+1} = cos(pos / 10000^{2i / d_model})
    '''
    def call(self, inputs):
        fl_type = inputs.dtype
        batch_size, max_length, depth = tf.unstack(tf.shape(inputs))

        depth_counter = tf.range(depth) // 2 * 2  # 0, 0, 2, 2, 4, ...
        depth_matrix = tf.tile(tf.expand_dims(depth_counter, 0), [max_length, 1])  # [max_length, depth]
        depth_matrix = tf.pow(10000.0, tf.cast(depth_matrix / depth, fl_type))  # [max_length, depth]

        # cos(x) == sin(x + π/2)
        phase = tf.cast(tf.range(depth) % 2, fl_type) * math.pi / 2  # 0, π/2, 0, π/2, ...
        phase_matrix = tf.tile(tf.expand_dims(phase, 0), [max_length, 1])  # [max_length, depth]

        pos_counter = tf.range(max_length)
        pos_matrix = tf.cast(tf.tile(tf.expand_dims(pos_counter, 1), [1, depth]), fl_type)  # [max_length, depth]

        positional_encoding = tf.sin(pos_matrix / depth_matrix + phase_matrix)
        # [batch_size, max_length, depth]
        positional_encoding = tf.tile(tf.expand_dims(positional_encoding, 0), [batch_size, 1, 1])

        return inputs + positional_encoding
```

PositionalEncodingの出力を確認します。

```python
import matplotlib.pyplot as plt

layer = AddPositionalEncoding()
test = tf.constant(np.zeros((32,128,512)), tf.float32)
layer(test).shape

plt.imshow(layer(test).numpy()[0]);
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2358656/df72b3d2-458d-9484-c984-3ccf9629f7ef.png)

Embedding レイヤの実装をします。
学習済みの重みを指定できるようにしましたが、今回は使用しません。

```python
class TokenEmbedding(Layer):
    def __init__(self, vocab_size, embedding_dim, embeddings=None, PAD_ID=0):
        # vocab_size: 単語の総数
        # embedding_dim: Embeddingの次数
        super().__init__()
        self.pad_id = PAD_ID
        self.embedding_dim = embedding_dim
        
        self.embedding = Embedding(vocab_size, embedding_dim)
        
        if embeddings is None:
            self.embedding = Embedding(input_dim=vocab_size,
                                       output_dim=embedding_dim,
                                       mask_zero=True,
                                       trainable=True)
        else:
            self.embedding = Embedding(input_dim=embeddings.shape[0],
                                       output_dim=embeddings.shape[1],
                                       mask_zero=True,
                                       trainable=True,
                                       weights=[embeddings])

    def call(self, x):
        embedding = self.embedding(x)
        return embedding * self.embedding_dim ** 0.5
```

TransformerBlockを実装します。
MultiHeadAttentionとFeedForwardNetworkから構成されます。

```python
class TransformerBlock(Layer):
    """
    transformer block : before ->[attention -> FF]-> next
    それぞれ残差接続とLayerNormalizationの処理が含まれる
    """
    def __init__(self, hidden_dim, heads_num, drop_rate=0.1):
        """
        hidden_numはheads_numで割り切れえる値とすること
        """
        super().__init__()
        self.atten = ResidualNormalizationWrapper(
            layer = MultiHeadAttention(hidden_dim = hidden_dim, heads_num = heads_num, drop_rate = drop_rate),
            drop_rate = drop_rate)
        
        self.ffn = ResidualNormalizationWrapper(
            layer = FeedForwardNetwork(hidden_dim = hidden_dim, drop_rate = drop_rate),
            drop_rate = drop_rate)
    
    def call(self, input, training, memory=None, attention_mask=None, return_attention_scores=False):
        """
        入力と出力で形式が変わらない
        [batch_size, token_num, hidden_dim]
        """
        
        if return_attention_scores:
            x, attn_weights = self.atten(input,training, memory, attention_mask, return_attention_scores)
            x = self.ffn(x)
            return x, attn_weights
        else:
            x = self.atten(input, training, memory, attention_mask, return_attention_scores)
            x = self.ffn(x)
            return x
```

Encoderの実装をします。
TransformerBlockを繰り替えし適用します。

```python
class Encoder(Layer):
    '''
    TransformerのEncoder
    '''
    def __init__(
            self,
            vocab_size, # 単語の総数
            hopping_num, # Multi-head Attentionの繰り返し数
            heads_num, # Multi-head Attentionのヘッド数
            hidden_dim, # Embeddingの次数
            token_num, # 系列長(文章中のトークン数)
            drop_rate, # ドロップアウトの確率
            embeddings=None
    ):
        super().__init__()
        self.hopping_num = hopping_num
        
        # Embedding層
        self.token_embedding = TokenEmbedding(vocab_size, hidden_dim, embeddings)
        # Position Embedding
        self.add_position_embedding = AddPositionalEncoding()
        self.input_dropout_layer = Dropout(drop_rate)

        # Multi-head Attentionの繰り返し(hopping)のリスト
        self.attention_block_list = [TransformerBlock(hidden_dim, heads_num) for _ in range(hopping_num)]
        self.output_normalization = LayerNormalization()

    def call(
            self,
            input,
            training,
            memory=None,
            attention_mask=None,
            return_attention_scores=False
    ):
        '''
        input: 入力 [batch_size, length]
        memory: 入力 [batch_size, length]
        attention_mask: attention weight に適用される mask
            [batch_size, 1, q_length, k_length] 
            pad 等無視する部分が 0 となるようなもの(Decoderで使用)
        return_attention_scores : attention weightを出力するか
        出力 [batch_size, length, hidden_dim]
        '''
        # [batch_size, token_num] -> [batch_size, token_num, hidden_dim]
        embedded_input = self.token_embedding(input)
        # Positional Embedding
        embedded_input = self.add_position_embedding(embedded_input)
        query = self.input_dropout_layer(embedded_input, training=training)
        
        
        if return_attention_scores:
            # MultiHead Attentionを繰り返し適用
            for i in range(self.hopping_num):
                query, atten_weights = self.attention_block_list[i](query, training, memory, attention_mask, return_attention_scores)

            # [batch_size, token_num, hidden_dim]
            return self.output_normalization(query), atten_weights
        else:
            # MultiHead Attentionを繰り返し適用
            for i in range(self.hopping_num):
                query = self.attention_block_list[i](query, training, memory, attention_mask, return_attention_scores)

            # [batch_size, token_num, hidden_dim]
            return self.output_normalization(query)
```

Encoderからの出力を全結合層に入力することで、分類モデルを作成します。

```python
class AttentionClassifier(Model):
    def __init__(
            self,
            vocab_size, # 単語の総数
            hopping_num, # Multi-head Attentionの繰り返し数
            heads_num, # Multi-head Attentionのヘッド数
            hidden_dim, # Embeddingの次数
            token_num, # 系列長(文章中のトークン数)
            drop_rate, # ドロップアウトの確率
            NUMLABELS, # クラス数
            embeddings = None,
            PAD_ID = 0
    ):
        super().__init__()
        self.PAD_ID = PAD_ID
        
        self.encoder = Encoder(vocab_size, hopping_num, heads_num, hidden_dim, token_num, drop_rate, embeddings)
        self.dense1 = Dense(hidden_dim, activation='tanh')
        self.dropout1 = Dropout(drop_rate)   
        self.final_layer = Dense(NUMLABELS, activation='softmax')

    def call(self, x, training, return_attention_scores=False):
        self_attention_mask=self._create_enc_attention_mask(x)
        
        # [batch_size, token_num] -> [batch_size, token_num, hidden_dim]
        if return_attention_scores:
            enc_output, atten_weights = self.encoder(x, attention_mask=self_attention_mask,return_attention_scores=return_attention_scores)
        else:
            enc_output = self.encoder(x, attention_mask=self_attention_mask,return_attention_scores=return_attention_scores)
        
        # 文頭の重みを使用 [batch_size, 0, hidden_dim]
        # [batch_size, hidden_dim] -> [batch_size, hidden_dim]
        enc_output = self.dense1(enc_output[:, 0, :])
        enc_output = self.dropout1(enc_output)
        
        # [batch_size, hidden_dim] -> [batch_size, NUMLABELS]
        final_output = self.final_layer(enc_output)

        if return_attention_scores:
            return final_output, atten_weights
        else:
            return final_output
    
    def _create_enc_attention_mask(self, x):
        batch_size, length = tf.unstack(tf.shape(x))
        # マスクする部分を1とする
        pad_array = tf.cast(tf.equal(x, self.PAD_ID), tf.float32)  # [batch_size, token_num]
        
        # shape broadcasting で [batch_size, head_num, token_num, token_num] になる
        return tf.reshape(pad_array, [batch_size, 1, 1, length])
```

データの準備をします。
自分で集めた居酒屋の口コミデータです。
まずはボキャブラリの関数を定義します。

```python
import re
import pandas as pd
from janome.tokenizer import Tokenizer
j_t = Tokenizer(wakati=True)

class Vocab(object):
    def __init__(self):
        self.w2i = {}
        self.i2w = {}
        self.special_chars = ['<pad>', '
', '
', '<unk>']
        self.bos_char = self.special_chars[1]
        self.eos_char = self.special_chars[2]
        self.oov_char = self.special_chars[3]

    def fit(self, sentences, path=None):
        self._words = set()

        #with open(path, 'r',encoding="utf-8") as f:
        #    sentences = f.read().splitlines()

        for sentence in sentences:
            #self._words.update(sentence.split())
            self._words.update(sentence)

        self.w2i = {w: (i + len(self.special_chars))
                    for i, w in enumerate(self._words)}

        for i, w in enumerate(self.special_chars):
            self.w2i[w] = i

        self.i2w = {i: w for w, i in self.w2i.items()}

    def transform(self, sentences, path=None, bos=False, eos=False):
        output = []

        #with open(path, 'r',encoding="utf-8") as f:
        #    sentences = f.read().splitlines()

        for sentence in sentences:
            #sentence = sentence.split()
            if bos:
                sentence = [self.bos_char] + sentence
            if eos:
                sentence = sentence + [self.eos_char]
            output.append(self.encode(sentence))

        return output

    def encode(self, sentence):
        output = []

        for w in sentence:
            if w not in self.w2i:
                idx = self.w2i[self.oov_char]
            else:
                idx = self.w2i[w]
            output.append(idx)

        return output

    def decode(self, sentence):
        return [self.i2w[id] for id in sentence]


def tokenizer_janome(text):
    return [tok for tok in j_t.tokenize(text, wakati=True)]

def preprocessing_text(text):
    text = re.sub('\r', '', text)
    text = re.sub('\n', '', text)
    text = re.sub('　', '', text)
    #text = re.sub(' ', '', text)
    
    text = re.sub(r'[0-9 ０-９]', '0', text)
    return text

def tokenizer_with_preprocessing(text):
    text = preprocessing_text(text)
    ret = tokenizer_janome(text)
    return ret
```

データの読み込みをします。

```python
from sklearn.model_selection import train_test_split

path='reviews.csv'

df = pd.read_csv(path)
seq_row = df['Body']
y = df['Rating2']

seq = [tokenizer_with_preprocessing(text) for text in seq_row]

x_train, x_test, y_train, y_test = train_test_split(seq, y, stratify=y)

vocab = Vocab()
vocab.fit(x_train)

x_id_train = vocab.transform(x_train, bos=True)
x_id_test = vocab.transform(x_test, bos=True)

X_train = pad_sequences(x_id_train, padding='post', maxlen=64)
y_train_oht = tf.one_hot(y_train, depth=2, dtype=tf.float32)
```

ネットワークの定義と、学習の設定をします。

```python
from tensorflow.keras import optimizers

model = AttentionClassifier(
            vocab_size = len(vocab.i2w), # 単語の総数
            hopping_num = 8, # Multi-head Attentionの繰り返し数
            heads_num = 6, # Multi-head Attentionのヘッド数
            hidden_dim = 300, # Embeddingの次数
            drop_rate = 0.1, # ドロップアウトの確率
            token_num = 64,
            NUMLABELS = 2
)

criterion = tf.keras.losses.CategoricalCrossentropy()
optimizer = optimizers.Adam(learning_rate=2e-4,
                           beta_1=0.9, beta_2=0.999, amsgrad=True)

model.compile(loss=criterion, optimizer=optimizer, metrics=['accuracy'])
```

学習を実行します。

```python
history=model.fit(X_train, y_train_oht, batch_size=32, epochs=20)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2358656/2007f7f2-a2a1-e482-b00c-c4e5fa471f45.png)

結果の可視化をします。

```python
from IPython.display import HTML
from IPython.display import HTML, display

def mk_html(index, batch, preds, normlized_weights, vocab, labels=["Negative", "Positive"]):
    "HTMLデータを作成する"

    # indexの結果を抽出
    sentence = batch[0][index]  # 文章
    label = batch[1][index] # ラベル
    pred = preds[index]  # 予測
    # ラベルと予測結果を文字に置き換え
    label_str = labels[label]
    pred_str = labels[pred]
    # 表示用のHTMLを作成する
    html = '正解ラベル：{}<br>推論ラベル：{}<br><br>'.format(label_str, pred_str)

    # 12種類のAttentionの平均を求める。最大値で規格化
    all_attens = normlized_weights[0, :, 0, :].sum(axis=0)*0  # all_attensという変数を作成する
    all_attens = np.sum(normlized_weights[index, :, 0, :], axis=0)
    all_attens = (all_attens -all_attens.min()) /  (all_attens.max()-all_attens.min())

    for word, attn in zip(sentence, all_attens):
        # 単語が[SEP]の場合は文章が終わりなのでbreak
        if vocab.decode([word.numpy().tolist()])[0] == "[SEP]":
            break

        # 関数highlightで色をつける、関数tokenizer_bert.convert_ids_to_tokensでIDを単語に戻す
        html += highlight(vocab.decode(
            [word.numpy().tolist()])[0], attn)
    html += "<br><br>"

    return html

def highlight(word, attn):
    "Attentionの値が大きいと文字の背景が濃い赤になるhtmlを出力させる関数"

    html_color = '#%02X%02X%02X' % (
        255, int(255*(1 - attn)), int(255*(1 - attn)))
    return '<span style="background-color: {}"> {}</span>'.format(html_color, word)
  
def make_html(preds, batch, labels=["Negative", "Positive"]):
  html_output = [mk_html(index=idx,
                         batch=batch, 
                         preds=np.argmax(preds, axis=1),
                         normlized_weights=atten.numpy(),
                         vocab=vocab,
                         labels=labels) for idx in np.arange(len(preds))]
  return  html_output
```

```python
preds, atten = model(X_test, return_attention_scores=True)
html_results = make_html(preds, batch=(X_test, np.array(y_test)), labels=["Negative", "Positive"])

HTML(html_results[613])
```

文章の単語ごとにattention weightで色付けしたものです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2358656/255c0a61-54a9-ae1d-f5c5-f8c272978129.png)

kerasではMultiHeadAttentionが実装されています。
引数名が変わることに注意して書き換えます。

```python
MultiHeadAttention(key_dim = 2, num_heads = heads_num, dropout = drop_rate)
```

他にも引数名の違いによる書き換えがあります。
詳細の実装・処理結果は以下を参考してください。

https://github.com/tanakakao/test/blob/969ae4f9e156ed0b23a267e48b5379ae75e85917/transformer_keras.ipynb


### pytorch
pytorchでも同様に実装をしていきます。
まずは必要なライブラリのインポートをします。


```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optimizers

import numpy as np
import math
```

kerasの実装をできるだけそのままpytorchに移し替えたので内容は変わりません。

MultiHeadAttentionの実装をします。

```python
class MultiHeadAttention(nn.Module):
    '''
    Multi-Head Attentionレイヤ
    
    hidden_dim : Embeddingされた単語ベクトルの長さ
    heads_num : マルチヘッドAttentionのヘッド数
       ※hidden_numはheads_numで割り切れえる値とすること
    drop_rate : 出力のDropout率

    model = MultiheadAttention(
        hidden_dim = 512,
        head_num = 8,
        drop_rate = 0.5
    )
    '''
    def __init__(self, token_num, hidden_dim, heads_num, drop_rate=0.5):
        super(MultiHeadAttention, self).__init__()
        # 入力の線形変換
        # 重み行列は[hidden_dim, hidden_dim]
        self.query = nn.Linear(hidden_dim, hidden_dim)
        self.key   = nn.Linear(hidden_dim, hidden_dim)
        self.value = nn.Linear(hidden_dim, hidden_dim)
        
        # 出力の線形変換
        self.projection = nn.Linear(hidden_dim, hidden_dim)
        
        # 出力のDropout
        self.drop = nn.Dropout(drop_rate)
        
        self.nf = hidden_dim
        self.nh = heads_num
    
    def atten(self, query, key, value, attention_mask):
        """
        Attention
        
        query, key, value : 入力
        attention_mask : attention weight に適用される mask
        """
        # 各値を取得
        shape = query.shape
        batch_size = -1 if shape[0] is None else shape[0]
        token_num = shape[2] # トークン列数
        hidden_dim = shape[1]*shape[3] # 入力チャンネル数
        
        # ここで q と k の内積を取ることで、query と key の単語間の関連度のようなものを計算します。
        # tf.matmulで最後の2成分について積を計算(それ以外は形がそろっている必要あり)
        # transpose_bで転置
        # [token_num, hidden_dim/head_num] @ [hidden_dim/head_num, token_num] = [token_num, token_num]
        scores = torch.matmul(query, key.transpose(-2, -1))
        
        # scoreをhidden_dimの平方根割る
        scores = scores / math.sqrt(hidden_dim)
        
        # Attention Maskがあればscoreに加算
        # attention_mask: [batch_size, token_num, token_num] 
        # マスク(参照しない部分)の場所に1、使用する部分は0とする
        # 0の部分を -無限大にする(softmax(-無限大)=0となる)
        # 1. PADを無視
        # 2. DecoderのSelf-Attentionで未来の情報を参照できないようにする
        if attention_mask is not None:
            scores = scores.masked_fill(attention_mask == 1, -1e9)

        # softmax を取ることで正規化します
        # input(query) の各単語に対して memory(key) の各単語のどこから情報を引いてくるかの重み
        atten_weight = F.softmax(scores, dim = -1)
        #atten_weight = scores / torch.sum(scores, dim=-1, keepdim=True)
        
        # 重みに従って value から情報を引いてきます
        # [token_num, token_num] @ [token_num, hidden_dim/head_num] = [token_num, hidden_dim/head_num]
        # input(query) の単語ごとに memory(value)の各単語 に attention_weight を掛け合わせて足し合わせた ベクトル(分散表現の重み付き和)を計算
        context = torch.matmul(atten_weight, value)
        
        # 各ヘッドの結合(reshape)
        # 入力と同じ形に変換する
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, token_num, hidden_dim)
        
        # 線形変換
        context = self.projection(context)
        
        return self.drop(context), atten_weight

    def _split(self, x):
        """
        query, key, valueを分割する
        
        入力 shape: [batch_size, length, hidden_dim] の時
        出力 shape: [batch_size, head_num, length, hidden_dim//head_num]
        """
        # 各値を取得
        hidden_dim = self.nf
        heads_num = self.nh
        shape = x.shape
        batch_size = -1 if shape[0] is None else shape[0]
        token_num = shape[1] # トークン列数
        
        # [batch_size, token_num, hidden_dim] -> [batch_size, token_num, head_num, hidden_dim/head_num]
        # splitだが実際は次元を拡張する処理
        x = x.view(batch_size, token_num, heads_num, int(hidden_dim/heads_num))
        
        # [batch_size, token_num, head_num, hidden_dim/head_num] -> [batch_size, head_num, token_num, hidden_dim/head_num]
        x = x.transpose(1, 2)
        return x
    
    def forward(self, x, memory=None, attention_mask=None, return_attention_scores=False):
        """
        モデルの実行
        
        input : 入力(query) [batch_size, token_num, hidden_dim]
        memory : 入力(key, value) [batch_size, token_num, hidden_dim]
        attention_mask : attention weight に適用される mask
            [batch_size, 1, q_length, k_length] 
            pad 等無視する部分が 1 となるようなもの(Decoderで使用)
        """
        # memoryが入力されない場合、memory=input(Self Attention)とする
        if memory is None:
            memory = x
        
        # input -> query
        # memory -> key, value
        # [batch_size, token_num, hidden_dim] @ [hidden_dim, hidden_dim] -> [batch_size, token_num, hidden_dim] 
        query = self.query(x)
        key = self.key(memory)
        value = self.value(memory)
        
        # ヘッド数に分割する
        # 実際はreshapeで次数を1つ増やす
        # [batch_size, token_num, hidden_dim] -> [batch_size, head_num, token_num, hidden_dim/head_num]
        query = self._split(query)
        key = self._split(key)
        value = self._split(value)
        
        # attention
        # 入力と同じ形の出力
        # context: [batch_size, token_num, hidden_dim]
        # score_weightsはEncoderではNoneとする
        context, atten_weight = self.atten(query, key, value, attention_mask)
        
        if return_attention_scores:
            return context, atten_weight
        else:
            return context

```

FeedForwardNetworkの実装をします。


```python
class FeedForwardNetwork(nn.Module):
    '''
    Position-wise Feedforward Neural Network
    transformer blockで使用される全結合層
    '''
    def __init__(self, hidden_dim, drop_rate=0.1):
        super().__init__()
        # 2層構造
        # 1層目：チャンネル数を増加させる
        self.filter_dense_layer = nn.Linear(hidden_dim, hidden_dim * 4)
        self.relu1 = nn.ReLU()
        
        # 2層目：元のチャンネル数に戻す
        self.output_dense_layer = nn.Linear(hidden_dim * 4, hidden_dim)
        self.drop = nn.Dropout(drop_rate)

    def forward(self, x):
        '''
        入力と出力で形が変わらない
        [batch_size, token_num, hidden_dim]
        '''
        
        # [batch_size, token_num, hidden_dim] -> [batch_size, token_num, 4*hidden_dim]
        x = self.filter_dense_layer(x)
        x = self.relu1(x)
        x = self.drop(x)
        
        # [batch_size, token_num, 4*hidden_dim] -> [batch_size, token_num, hidden_dim]
        return self.output_dense_layer(x)
```

残差接続のためのResidualNormalizationWrapperを実装します。

```python
class ResidualNormalizationWrapper(nn.Module):
    '''
    残差接続
    output: input + SubLayer(input)
    '''
    def __init__(self, hidden_dim, layer, drop_rate=0.1):
        super().__init__()
        self.layer = layer # SubLayer : ここではAttentionかFFN
        self.layer_normalization = nn.LayerNorm(hidden_dim)
        self.drop = nn.Dropout(drop_rate)

    def forward(self, x, memory=None, attention_mask=None, return_attention_scores=False):
        """
        AttentionもFFNも入力と出力で形が変わらない
        [batch_size, token_num, hidden_dim]
        """
        
        params = {}
        if memory is not None:
            params['memory'] = memory
        if attention_mask is not None:
            params['attention_mask'] = attention_mask
        if return_attention_scores:
            params['return_attention_scores'] = return_attention_scores
        
        out = self.layer_normalization(x)
        if return_attention_scores:
            out, attn_weights = self.layer(out,**params)
            out = self.drop(out)
            return x + out, attn_weights
        else:
            out = self.layer(out,**params)
            out = self.drop(out)
            return x + out
```

AddPositionalEncodingを実装します。
これはいろいろ実装の種類があるようです。

```python
class AddPositionalEncoding(nn.Module):
    '''
    入力テンソルに対し、位置の情報を付与して返すレイヤー
    see: https://arxiv.org/pdf/1706.03762.pdf

    PE_{pos, 2i}   = sin(pos / 10000^{2i / d_model})
    PE_{pos, 2i+1} = cos(pos / 10000^{2i / d_model})
    '''
    def forward(self, inputs):
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        fl_type = inputs.dtype
        batch_size, max_length, depth = inputs.shape
        
        depth_counter = torch.div(torch.arange(depth) ,2, rounding_mode='trunc')*2
 
        depth_matrix = torch.tile(torch.unsqueeze(depth_counter, 0), [max_length, 1])  # [max_length, depth]
        depth_matrix = torch.pow(10000.0, depth_matrix / depth)  # [max_length, depth]
        # cos(x) == sin(x + π/2)
        phase = torch.remainder(torch.arange(depth), 2) * math.pi / 2
        phase_matrix = torch.tile(torch.unsqueeze(phase, 0), [max_length, 1])  # [max_length, depth]

        pos_counter = torch.arange(max_length)
        pos_matrix = (torch.tile(torch.unsqueeze(pos_counter, 1), [1, depth]))  # [max_length, depth]

        positional_encoding = torch.sin(pos_matrix / depth_matrix + phase_matrix)
        # [batch_size, max_length, depth]
        positional_encoding = torch.tile(torch.unsqueeze(positional_encoding, 0), [batch_size, 1, 1])
        positional_encoding = positional_encoding.to(device)

        return inputs + positional_encoding
```

TokenEmbeddingを実装します。
学習済みの重みを使えるようにしていますが、今回は試していません。

```python
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, embedding_dim, pretrained_weight=None):
        # vocab_size: 単語の総数
        # embedding_dim: Embeddingの次数
        super().__init__()
        self.embedding_dim = embedding_dim
        
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=1)
        self.embedding.weight.requires_grad = True
        
        if pretrained_weight is not None:
            self.embedding.weight.data.copy_(pretrained_weight)

    def forward(self, x):
        # inputのIDに対応したベクトルを持ってくる
        embedding = self.embedding(x)
        
        return embedding * (self.embedding_dim ** 0.5)
```

TransformerBlockを実装します。

```python
class TransformerBlock(nn.Module):
    """
    transformer block : before ->[attention -> FF]-> next
    それぞれ残差接続とLayerNormalizationの処理が含まれる
    """
    def __init__(self, token_num, hidden_dim, heads_num, drop_rate=0.1):
        """
        hidden_numはheads_numで割り切れえる値とすること
        """
        super().__init__()
        self.atten = ResidualNormalizationWrapper(
            hidden_dim = hidden_dim,
            layer = MultiHeadAttention(token_num=token_num, hidden_dim = hidden_dim, heads_num = heads_num, drop_rate = drop_rate),
            drop_rate = drop_rate)
        
        self.ffn = ResidualNormalizationWrapper(
            hidden_dim = hidden_dim,
            layer = FeedForwardNetwork(hidden_dim = hidden_dim, drop_rate = drop_rate),
            drop_rate = drop_rate)
    
    def forward(self, input, memory=None, attention_mask=None, return_attention_scores=False):
        """
        入力と出力で形式が変わらない
        [batch_size, token_num, hidden_dim]
        """
        if return_attention_scores:
            x, attn_weights = self.atten(input, memory, attention_mask, return_attention_scores)
            x = self.ffn(x)
            return x, attn_weights
        else:
            x = self.atten(input, memory, attention_mask, return_attention_scores)
            x = self.ffn(x)
            return x
```

TransformerBlockを繰り返し実行してEncoderを定義します。


```python
class Encoder(nn.Module):
    '''
    TransformerのEncoder
    '''
    def __init__(
            self,
            vocab_size, # 単語の総数
            hopping_num, # Multi-head Attentionの繰り返し数
            heads_num, # Multi-head Attentionのヘッド数
            hidden_dim, # Embeddingの次数
            token_num, # 系列長(文章中のトークン数)
            drop_rate, # ドロップアウトの確率
            pretrained_weight=None
    ):
        super().__init__()
        self.hopping_num = hopping_num
        
        # Embedding層
        self.token_embedding = TokenEmbedding(vocab_size, hidden_dim, pretrained_weight)
        # Position Embedding
        self.add_position_embedding = AddPositionalEncoding()
        self.input_dropout_layer = nn.Dropout(drop_rate)

        # Multi-head Attentionの繰り返し(hopping)のリスト
        self.attention_block_list = nn.ModuleList([TransformerBlock(token_num, hidden_dim, heads_num) for _ in range(hopping_num)])
        self.output_normalization = nn.LayerNorm(hidden_dim)

    def forward(
            self,
            input,
            memory=None,
            attention_mask=None,
            return_attention_scores=False
    ):
        '''
        input: 入力 [batch_size, length]
        memory: 入力 [batch_size, length]
        attention_mask: attention weight に適用される mask
            [batch_size, 1, q_length, k_length] 
            pad 等無視する部分が 0 となるようなもの(Decoderで使用)
        出力 [batch_size, length, hidden_dim]
        '''
        # [batch_size, token_num] -> [batch_size, token_num, hidden_dim]
        embedded_input = self.token_embedding(input)
        # Positional Embedding
        embedded_input = self.add_position_embedding(embedded_input)
        query = self.input_dropout_layer(embedded_input)
        
        if return_attention_scores:
            # MultiHead Attentionを繰り返し適用
            for i in range(self.hopping_num):
                query, atten_weights = self.attention_block_list[i](query, memory, attention_mask, return_attention_scores)

            # [batch_size, token_num, hidden_dim]
            return self.output_normalization(query), atten_weights
        else:
            # MultiHead Attentionを繰り返し適用
            for i in range(self.hopping_num):
                query = self.attention_block_list[i](query, memory, attention_mask, return_attention_scores)

            # [batch_size, token_num, hidden_dim]
            return self.output_normalization(query)
```

TransformerBlockの出力を全結合層に入力して分類モデルとします。

```python
class AttentionClassifier(nn.Module):
    def __init__(
            self,
            vocab_size, # 単語の総数
            hopping_num, # Multi-head Attentionの繰り返し数
            heads_num, # Multi-head Attentionのヘッド数
            hidden_dim, # Embeddingの次数
            token_num, # 系列長(文章中のトークン数)
            drop_rate, # ドロップアウトの確率
            NUMLABELS, # クラス数
            pretrained_weight=None,
            PAD_ID = 1
    ):
        super().__init__()
        self.PAD_ID = PAD_ID
        
        self.encoder = Encoder(vocab_size, hopping_num, heads_num, hidden_dim, token_num, drop_rate, pretrained_weight)
        self.dense1 = nn.Linear(hidden_dim, hidden_dim)
        self.act1 = nn.Tanh()
        self.dropout1 = nn.Dropout(drop_rate)   
        self.final_layer = nn.Linear(hidden_dim, NUMLABELS)
        
        nn.init.normal_(self.dense1.weight, std=0.02)
        nn.init.normal_(self.dense1.bias, std=0)
        nn.init.normal_(self.final_layer.weight, std=0.02)
        nn.init.normal_(self.final_layer.bias, std=0)

    def forward(self, x, return_attention_scores=False):
        self_attention_mask=self._create_enc_attention_mask(x)
        
        # [batch_size, token_num] -> [batch_size, token_num, hidden_dim]
        if return_attention_scores:
            enc_output, atten_weights = self.encoder(x, attention_mask=self_attention_mask,return_attention_scores=return_attention_scores)
        else:
            enc_output = self.encoder(x, attention_mask=self_attention_mask,return_attention_scores=return_attention_scores)
        
        # 文頭の重みを使用 [batch_size, 0, hidden_dim]
        # [batch_size, hidden_dim] -> [batch_size, hidden_dim]
        enc_output = self.dense1(enc_output[:, 0, :])
        enc_output = self.act1(enc_output)
        enc_output = self.dropout1(enc_output)
        
        # [batch_size, hidden_dim] -> [batch_size, NUMLABELS]
        final_output = self.final_layer(enc_output)

        if return_attention_scores:
            return final_output, atten_weights
        else:
            return final_output
    
    def _create_enc_attention_mask(self, x):
        batch_size, length = x.shape
        # マスクする部分を1とする
        pad_array = torch.eq(x, self.PAD_ID).to(dtype=torch.int8)  # [batch_size, token_num]
        
        # shape broadcasting で [batch_size, head_num, token_num, token_num] になる
        return pad_array.view([batch_size, 1, 1, length])
```

データの準備を行います。
自分で集めた居酒屋の口コミデータです。
まずは、前処理を定義します。
今回は、簡単な処理のみとします。

```python
!pip install janome
import re
from janome.tokenizer import Tokenizer
j_t = Tokenizer(wakati=True)

def tokenizer_janome(text):
    return [tok for tok in j_t.tokenize(text, wakati=True)]

def preprocessing_text(text):
    text = re.sub('\r', '', text)
    text = re.sub('\n', '', text)
    text = re.sub('　', '', text)
    text = re.sub(' ', '', text)
    
    text = re.sub(r'[0-9 ０-９]', '0', text)
    return text

def tokenizer_with_preprocessing(text):
    text = preprocessing_text(text)
    ret = tokenizer_janome(text)
    return ret
```

データの読み込みをします。


```python
import torchtext
#from torchtext import data, datasets
from torchtext.legacy import data

max_length = 64
TEXT = data.Field(sequential=True, tokenize=tokenizer_with_preprocessing,
                  use_vocab=True, lower=True, include_lengths=True,
                  batch_first=True, fix_length=max_length,init_token="<eos>",eos_token="<cls>")
LABEL = data.Field(sequential=False, use_vocab=False, preprocessing=None)

dataset = data.TabularDataset(
        path='reviews.csv', format='csv',
        skip_header=True,
        fields=[('Text', TEXT), ('Label', LABEL), ('Label2', LABEL)])

train_dataset, test_dataset = dataset.split(split_ratio=0.7)
train_dataset, val_dataset = train_dataset.split(split_ratio=0.7)
```

Iteratorの定義をします。

```python
train_iter = data.Iterator(
    train_dataset, batch_size=32, 
    train=True  # train=Trueならシャッフルソートは有効
)
val_iter = data.Iterator(
    val_dataset, batch_size=32, 
    train=False, sort=False
)
test_iter = data.Iterator(
    test_dataset, batch_size=32, 
    train=False, sort=False
)

dataloaders_dict = {"train":train_iter, "val":val_iter}
```

ネットワークの定義や学習の設定をします。

```python
TEXT.build_vocab(train_dataset)
vocab = TEXT.vocab

net = AttentionClassifier(
            vocab_size = len(vocab), # 単語の総数
            hopping_num = 8, # Multi-head Attentionの繰り返し数
            heads_num = 6, # Multi-head Attentionのヘッド数
            hidden_dim = 300, # Embeddingの次数
            drop_rate = 0.1, # ドロップアウトの確率
            token_num = 64,
            
    pretrained_weight=None,
    NUMLABELS=2
    )

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.kaiming_normal_(m.weight)
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.0)
net.train()
net.apply(weights_init)

criterion = nn.CrossEntropyLoss()

learning_rate = 2e-4
optimizer = optimizers.Adam(net.parameters(), lr=learning_rate, amsgrad=True, eps=1e-07)
```

訓練用の関数を定義します。

```python
def train_model(net, dataloaders_dict, criterion, optimizer, num_epochs):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("使用デバイス:", device)
    print('--------start--------')
    net.to(device)
    
    torch.backends.cudnn.benchmark = True
    
    for epoch in range(num_epochs):
        for phase in ['train', 'val']:
            if phase == 'train':
                net.train()
            else:
                net.eval()
            
            epoch_loss = 0.0
            epoch_corrects = 0
            
            for batch in (dataloaders_dict[phase]):
                inputs = batch.Text[0].to(device)
                labels = batch.Label2.to(device)
                
                optimizer.zero_grad()
                
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = net(inputs)
                    loss = criterion(outputs, labels)
                    
                    _, preds = torch.max(outputs, 1)
                    
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
                    
                    epoch_loss += loss.item() * inputs.size(0)
                    epoch_corrects += torch.sum(preds == labels.data)
            
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double() / len(dataloaders_dict[phase].dataset)
            
            print('Epoch {}/{} | {:.^5} | Loss: {:.4f} Acc: {:.4f}'.format(epoch+1,
                                                                           num_epochs,
                                                                           phase,
                                                                           epoch_loss,
                                                                           epoch_acc))
        
    return net
```

訓練を実行します。

```python
num_epochs = 20

net_trained = train_model(net, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2358656/387e4d2b-9887-3aeb-9e95-8d8232472443.png)


結果の可視化を行います。


```python
from IPython.display import HTML
from IPython.display import HTML, display

def mk_html(index, batch, preds, normlized_weights, vocab, labels=["Negative", "Positive"]):
    "HTMLデータを作成する"

    # indexの結果を抽出
    sentence = batch[0][index]  # 文章
    label = batch[1][index] # ラベル
    pred = preds[index]  # 予測
    # ラベルと予測結果を文字に置き換え
    label_str = labels[label]
    pred_str = labels[pred]
    # 表示用のHTMLを作成する
    html = '正解ラベル：{}<br>推論ラベル：{}<br><br>'.format(label_str, pred_str)

    # 12種類のAttentionの平均を求める。最大値で規格化
    all_attens = normlized_weights[0, :, 0, :].sum(axis=0)*0  # all_attensという変数を作成する
    all_attens = np.sum(normlized_weights[index, :, 0, :], axis=0)
    all_attens = (all_attens -all_attens.min()) /  (all_attens.max()-all_attens.min())

    for word, attn in zip(sentence, all_attens):
        # 単語が[SEP]の場合は文章が終わりなのでbreak
        if TEXT.vocab.itos[word] == "<cls>":
            break

        # 関数highlightで色をつける、関数tokenizer_bert.convert_ids_to_tokensでIDを単語に戻す
        html += highlight(TEXT.vocab.itos[word], attn)
    html += "<br><br>"

    return html

def highlight(word, attn):
    "Attentionの値が大きいと文字の背景が濃い赤になるhtmlを出力させる関数"

    html_color = '#%02X%02X%02X' % (
        255, int(255*(1 - attn)), int(255*(1 - attn)))
    return '<span style="background-color: {}"> {}</span>'.format(html_color, word)
  
def make_html(preds, batch, labels=["Negative", "Positive"]):
  html_output = [mk_html(index=idx,
                         batch=batch, 
                         preds=np.argmax(preds, axis=1),
                         normlized_weights=atten,
                         vocab=vocab,
                         labels=labels) for idx in np.arange(len(preds))]
  return  html_output
```

```python
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
batch = next(iter(test_iter))

preds, atten = net(batch.Text[0].to(device), return_attention_scores=True)
atten = atten.to('cpu').detach().numpy()
preds = preds.to('cpu').detach().numpy()

html_results = make_html(preds, batch=(np.array(batch.Text[0]), np.array(batch.Label2)), labels=["Negative", "Positive"])
HTML(html_results[1])
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2358656/28631f78-c469-85b4-0e3f-4dae7089123a.png)


pytorchではnn.TransformerEncoderLayerおよびnn.TransformerEncoderが実装されており、これを使えばよりシンプルに実装ができます。
書き換えるのはEncoder部分です。

```python
class Encoder(nn.Module):
    '''
    TransformerのEncoder
    '''
    def __init__(
            self,
            vocab_size, # 単語の総数
            hopping_num, # Multi-head Attentionの繰り返し数
            heads_num, # Multi-head Attentionのヘッド数
            hidden_dim, # Embeddingの次数
            token_num, # 系列長(文章中のトークン数)
            drop_rate, # ドロップアウトの確率
            pretrained_weight=None
    ):
        super().__init__()
        self.hopping_num = hopping_num
        
        # Embedding層
        self.token_embedding = TokenEmbedding(vocab_size, hidden_dim, pretrained_weight)
        # Position Embedding
        self.add_position_embedding = AddPositionalEncoding()
        self.input_dropout_layer = nn.Dropout(drop_rate)

        # Multi-head Attentionの繰り返し(hopping)のリスト
        encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=heads_num, dropout=drop_rate, batch_first=True)
        self.attention_block = nn.TransformerEncoder(encoder_layer, num_layers=hopping_num)
        
        self.output_normalization = nn.LayerNorm(hidden_dim)

    def forward(
            self,
            input,
            memory=None,
            attention_mask=None
    ):
        '''
        input: 入力 [batch_size, length]
        memory: 入力 [batch_size, length]
        attention_mask: attention weight に適用される mask
            [batch_size, 1, q_length, k_length] 
            pad 等無視する部分が 0 となるようなもの(Decoderで使用)
        出力 [batch_size, length, hidden_dim]
        '''
        # [batch_size, token_num] -> [batch_size, token_num, hidden_dim]
        embedded_input = self.token_embedding(input)
        # Positional Embedding
        embedded_input = self.add_position_embedding(embedded_input)
        query = self.input_dropout_layer(embedded_input)
        
        query = self.attention_block(query, src_key_padding_mask=attention_mask)

        # [batch_size, token_num, hidden_dim]
        return self.output_normalization(query)
```

実装・処理結果は以下を参考してください。

https://github.com/tanakakao/test/blob/969ae4f9e156ed0b23a267e48b5379ae75e85917/transformer_torch2.ipynb

以上で、Transformer(Encoderのみ)の実装を終わります。