Transformerで日本語→英語の翻訳モデルを実装してみた

BERTの実装を試みていたが、先ずは前提となるTransformerを抑えておく必要があるため、この機会に日本語⇒英語の翻訳モデルを実装してみた。

NLPのタスクにおいて、これまでのRNN/LSTM/CNNに比べ、精度/計算コスト共に良い成果が得られたことから注目されているモデルである。
![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/398780/e6a6bbbe-1e80-c3fe-8088-4de0590b26fe.jpeg)



Transformerに関する論文に関しては、以下を参照。
**Attention Is All You Need:**
https://arxiv.org/abs/1706.03762
**Transformer: A Novel Neural Network Architecture for Language Understanding:**
https://research.googleblog.com/2017/08/transformer-novel-neural-network.html

<div align="center">
    <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/398780/c17b53b3-a16a-0bcd-6fd0-a35744963771.jpeg" alt="Transformer Model">
</div>


TransformerにおいてはAttention層が最も重要な役割を担っており、当該層が文章ごとにおける各単語間の関連度合いを算出し、かつ並列処理により計算速度、メモリ使用量の削減が可能とされる。
Attentionに関しては、以下記事において詳細に解説がされているため、本投稿では割愛する。

**作って理解する Transformer / Attention:**
https://qiita.com/halhorn/items/c91497522be27bde17ce

***◆Scaled Dot-Product Attention***
TransformerのベースであるMultiheadAttentionには、Scaled Dot-Product Attentionが使用されており、数式としては以下の通りである。
![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/398780/3b1d33f8-498b-9016-3909-badcb81cef72.jpeg)

入力データをK(key), V(Value), 出力データをQ(query)としており、当該値はベクトル化された単語である。
この中で重要なのは、Q, Kであり、当該各(ベクトル化された)単語の内積(matmul(q, k))により単語間の関連度合いを算出し、その値を√dkで除算することでスケールさせている。スケーリング因子を適用している理由としては、query, keyの次元数が大きい場合, 内積(matmul(q, k))の値が極端に大きくなり、バックプロパゲーションの際の勾配がそれに伴い縮小することで学習が進まなくなることを防ぐためである。



***◆MultiheadAttention***
Transformerでは、上記で述べたquery, key,valueを分割(マルチヘッド化)したAttentionが使用されている。
論文によると、単一ではなく分割した方が精度が上がったとの内容が記載されている。query, key,valueを分割処理させることで、分割された集合(ヘッド)が異なる位置の異なる部分空間を処理すると解釈でき，単一の場合は表現力の幅を狭めてしまうとのこと。
各ヘッドは出力時に連結(concat)され、最終的に順伝番層へ渡され出力される。
![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/398780/9c151c8d-abf1-d297-8b86-0ef2afd47a5e.jpeg)

Transformerの理論説明については、数多く文献があるので、以降は実際にPytorchを使用したTransformerの実装について記載する。

Transformerの実装に必要なクラスに関するドキュメントは以下公式ページを参照。
**Transformer Layers :**
https://pytorch.org/docs/stable/nn.html#transformer-layers







***1. 学習データの選定***
学習データは以下サイトのコーパスを使用。
データの中身としては、日本語⇔英語の対訳データが約300万件含まれている。
**JESC: Japanese-English Subtitle Corpus**
https://nlp.stanford.edu/projects/jesc/

***2. 必要ライブラリをインポート***

```python
import os
import sys
import re
import tarfile
import time
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from typing import Optional
```

***3. Pytorchモジュールをインポート***

```python
import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import torch.nn as nn
from torch.nn import Embedding, TransformerDecoderLayer, TransformerDecoder, TransformerEncoderLayer, TransformerEncoder
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torchtext.data import Field, BucketIterator, TabularDataset
from torch.nn.init import xavier_uniform_
```



***4. 実行環境をGPUに指定***

```python
device = "cuda" if torch.cuda.is_available() else "cpu"
```

***5. 学習用データをDataFrameに変換***

```python
path = './data/raw.tar.gz'
savedir = './data/'
with tarfile.open(path, 'r:*') as tar:
        tar.extractall(savedir)

filename = 'raw'
with open(os.path.join(savedir + 'raw/', filename), 'r') as f:
  raw_data = f.readlines()

raw_list = [re.sub('\n', '', s).split('\t') for s in raw_data]
raw_df = pd.DataFrame(raw_list,
                  columns=['英語', '日本語'])
```

**5.1 データの中身を確認**

```python
raw_df.head()
```

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/398780/e4badda2-dfcb-2eb9-dbc0-546e8f821955.jpeg)



***6. spaCyインストール***
Transformerの埋め込み層に入力する値に各単語のIDを入力する必要がある。各単語に紐づく単語IDリストを作成するために使用。

```python
!pip install spacy==2.3.4
import spacy
```

**7. 日本語・英語用のモデルをダウンロード**

```python
!python3 -m spacy download ja_core_news_md
!python3 -m spacy download en_core_web_md
```

**8. モデルの読み込み**

```python
JA = spacy.load("ja_core_news_md")
EN = spacy.load("en_core_web_md")
```

***9. 日本語・英語用Tokenizerを作成***

```python
def tokenize_ja(sentence):
    return [tok.text for tok in JA.tokenizer(sentence)]

def tokenize_en(sentence):
    return [tok.text for tok in EN.tokenizer(sentence)]
```

***10. Fieldオブジェクト定義***
Fieldを定義することで数値表現(Tensor)に変換することが可能

```python
JA_TEXT = Field(sequential=True, tokenize=tokenize_ja)
EN_TEXT = Field(sequential=True,
                tokenize=tokenize_en,
                init_token='<sos>',
                eos_token='<eos>',
                lower=True)
```

**11. 訓練/テスト用データに分割**

```python
train_val, test_val = train_test_split(raw_df, test_size=0.3)
```

***12. CSVファイルに書き出し***

```python
train_val.to_csv('train.csv', index = False)
test_val.to_csv('test.csv', index = False)
```

***13. 訓練・テスト用データのデータセット読み込み***

```python
train, test = TabularDataset.splits(
    path='./',
    train='train.csv',
    test='test.csv',
    format='csv',
    fields=[('en_text', EN_TEXT), ('ja_text', JA_TEXT)]
)
```

***14. Vocabオブジェクト作成***
これにより単語インデックス辞書/出現頻度/単語リスト、等が作成される。

```python
JA_TEXT.build_vocab(train, test)
EN_TEXT.build_vocab(train, test)
```

**14.1 日本語・英語双方の各単語にIDが付与されていることを確認**

```python
JA_TEXT.vocab.stoi
defaultdict(<function torchtext.vocab._default_unk_index>,
            {'<unk>': 0,
             '<pad>': 1,
             'の': 2,
             'は': 3,
             'に': 4,
             'て': 5,
             'た': 6,
             'を': 7,
             'が': 8,
             '.': 9,
             'だ': 10,
```
```python
EN_TEXT.vocab.stoi
defaultdict(<function torchtext.vocab._default_unk_index>,
            {'<unk>': 0,
             '<pad>': 1,
             '<sos>': 2,
             '<eos>': 3,
             '.': 4,
             ',': 5,
             'you': 6,
             'the': 7,
             'i': 8,
             '?': 9,
             'to': 10,
```

***15. 学習用バッチをイテレータオブジェクトとして作成***

```python
batch_size = 1024
train_iter, val_iter = BucketIterator.splits(
        datasets=(train, test),
        batch_size=batch_size,
        sort_key=lambda x: len(x.en_text)
        )

```

**15.1 バッチの中身を確認**

```python
batch = next(iter(train_iter))
tensor([[1902,  544, 1438,  ...,   55, 1181,  877],
        [   8,    7, 2202,  ...,    4,    7,    1],
        [ 671,  155,    1,  ...,  164,  171,    1],
        ...,
        [   1,    1,    1,  ...,    1,    1,    1],
        [   1,    1,    1,  ...,    1,    1,    1],
        [   1,    1,    1,  ...,    1,    1,    1]])
```

***16. 位置エンコーディング***
LSTM/RNN/CNNとは異なり、単語の位置関係を学習することができないため、位置関係を学習するために必要となる。
偶数インデックスにはsin、奇数インデックスにはcosを適用することで、単語間の依存関係を学習することが可能となり、文章内における単語の位置関係を把握することができる。
※ posは何番目の単語か、iは埋め込み層による分散表現の何番目の次元かを表している。
![https___imgur.com_xl12ivl.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/398780/0b4098cd-6789-d1b4-2092-ba96e4eb6224.png)


```python
class PositionalEncoding(nn.Module):
 
    def __init__(self, d_model, dropout: float = 0.1, max_len: int = 5000) -> None:
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model
 
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)
 
    def forward(self, x) -> torch.Tensor:
        x = x * math.sqrt(self.d_model)
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)
```

***17. Transformerモデル***

```python
class Transformer(nn.Module):
 
    def __init__(self, 
                 src_vocab_size: int = 3000,
                 d_model: int = 512,
                 nhead: int = 8,
                 dim_feedforward: int = 2048,
                 dropout: float = 0.5,
                 activation: str = 'relu',
                 num_layers: int = 6,
                 tgt_vocab_size: int = 3000
                 ) -> None:
        super(Transformer, self).__init__()
        self.src_embedd = nn.Embedding(src_vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model)

        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
        encoder_norm = nn.LayerNorm(d_model)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers, norm=encoder_norm)

        self.tgt_embedd = nn.Embedding(tgt_vocab_size, d_model)
        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout=dropout, activation=activation)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers, norm=decoder_norm)

        self.out = nn.Linear(d_model, tgt_vocab_size)
        self._reset_parameters()
        self.d_model = d_model
        self.nhead = nhead

    def _reset_parameters(self) -> None:
        for p in self.parameters():
            if p.dim() > 1:
                xavier_uniform_(p)

    def forward(self, 
                src: torch.Tensor,
                tgt: torch.Tensor,
                src_mask: Optional[torch.Tensor] = None,
                tgt_mask: Optional[torch.Tensor] = None,
                memory_mask: Optional[torch.Tensor] = None,
                src_key_padding_mask: Optional[torch.Tensor] = None,
                tgt_key_padding_mask: Optional[torch.Tensor] = None,
                memory_key_padding_mask: Optional[torch.Tensor] = None
                ) -> torch.Tensor:

      src = self.src_embedd(src)
      src = self.pos_encoder(src)
      memory = self.encoder(src)

      tgt = self.tgt_embedd(tgt)
      tgt = self.pos_encoder(tgt)
      out = self.decoder(tgt,
                          memory,
                          tgt_mask=tgt_mask,
                          memory_mask=memory_mask,
                          tgt_key_padding_mask=tgt_key_padding_mask,
                          memory_key_padding_mask=memory_key_padding_mask
                          )
      out = self.out(out)

      return out
```

***18. マスキング***
デコード層が予測すべき出力単語の情報が予測前のデコーダに知られた場合、予め正解を教えることとなり、学習が進まず汎化性能が下がるため、デコード層へ渡す文章に含まれる単語の一部にマスキングを実施する必要がある。
マスクには、（負の無限大に近い）-infを掛けることで、マスクがQとKのスケール済み行列積と合計され、softmax関数の直前に適用されることで、0に限りなく近い出力とするためである。

```python
def generate_square_subsequent_mask(sz: int) -> torch.Tensor:
  mask = (sz != 0)
  mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
  mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
  return mask
```

***19. モデルの初期化***

```python
src_vocab_size=len(JA_TEXT.vocab)
tgt_vocab_size=len(EN_TEXT.vocab)

model = Transformer(src_vocab_size=src_vocab_size, tgt_vocab_size=tgt_vocab_size)
model.to(device)
```

***20. 損失関数/最適化関数を定義***

```python
criterion = nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(), lr=0.0001)
```

***21. 学習***

```python
def train(epochs: int = 300) -> None:
  
  model.train()
  for epoch in range(epochs):
    total_loss = 0
    start_time = time.time()
    for i, data in enumerate(train_iter):

      src = data.ja_text
      tgt = data.en_text

      src = src.transpose(0, 1).to(device)
      tgt = tgt.transpose(0, 1)

      tgt_input = tgt[:, :-1].to(device)
      targets = tgt_input.contiguous().view(-1)

      tgt_mask = generate_square_subsequent_mask(tgt_input.size(1)).to(device)

      #学習/予測
      pred = model(src, tgt_input.transpose(0, 1), tgt_mask=tgt_mask)
      pred = pred.contiguous().view(-1, pred.size(-1))

      #損失計算
      loss = criterion(pred,targets)
      
      #バックプロパゲーション
      optim.zero_grad()
      loss.backward()
      optim.step()

      #サンプル単位の損失計算
      total_loss += loss.item() / batch_size

    end_time = time.time()
    if (epoch + 1) % 1 == 0:
      print('epoch: {}, total loss: {:.5g}, elasped time: {:.5g}'.format(epoch+1, total_loss, end_time-start_time))
```

***22. 予測***

```python
def predict(ja_src: str) -> str:

  src_list = [JA_TEXT.vocab.stoi[token] for token in tokenize_ja(ja_src)]

  src = torch.autograd.Variable(torch.LongTensor(src_list))
  src = src.unsqueeze(0)

  tgt = torch.zeros(src.size(1), src.size(0), dtype=torch.int64, requires_grad=False)

  tgt[0] = EN_TEXT.vocab.stoi['<sos>']

  src = src.transpose(0, 1).to(device)
  tgt = tgt.transpose(0, 1)

  tgt_input = tgt[:, :-1].to(device)

  tgt_mask = generate_square_subsequent_mask(tgt_input.size(1)).to(device)

  out = model(src, tgt_input.transpose(0, 1), tgt_mask)

  _, idx = torch.max(out[:, -1], 1)

  return " ".join([EN_TEXT.vocab.itos[i] for i in idx])
```