BERTの実装を試みていたが、先ずは前提となるTransformerを抑えておく必要があるため、この機会に日本語⇒英語の翻訳モデルを実装してみた。
NLPのタスクにおいて、これまでのRNN/LSTM/CNNに比べ、精度/計算コスト共に良い成果が得られたことから注目されているモデルである。
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
TransformerにおいてはAttention層が最も重要な役割を担っており、当該層が文章ごとにおける各単語間の関連度合いを算出し、かつ並列処理により計算速度、メモリ使用量の削減が可能とされる。
Attentionに関しては、以下記事において詳細に解説がされているため、本投稿では割愛する。
作って理解する Transformer / Attention:
https://qiita.com/halhorn/items/c91497522be27bde17ce
◆Scaled Dot-Product Attention
TransformerのベースであるMultiheadAttentionには、Scaled Dot-Product Attentionが使用されており、数式としては以下の通りである。
入力データを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)され、最終的に順伝番層へ渡され出力される。
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. 必要ライブラリをインポート
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モジュールをインポート
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に指定
device = "cuda" if torch.cuda.is_available() else "cpu"
5. 学習用データをDataFrameに変換
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 データの中身を確認
raw_df.head()
6. spaCyインストール
Transformerの埋め込み層に入力する値に各単語のIDを入力する必要がある。各単語に紐づく単語IDリストを作成するために使用。
!pip install spacy==2.3.4
import spacy
7. 日本語・英語用のモデルをダウンロード
!python3 -m spacy download ja_core_news_md
!python3 -m spacy download en_core_web_md
8. モデルの読み込み
JA = spacy.load("ja_core_news_md")
EN = spacy.load("en_core_web_md")
9. 日本語・英語用Tokenizerを作成
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)に変換することが可能
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. 訓練/テスト用データに分割
train_val, test_val = train_test_split(raw_df, test_size=0.3)
12. CSVファイルに書き出し
train_val.to_csv('train.csv', index = False)
test_val.to_csv('test.csv', index = False)
13. 訓練・テスト用データのデータセット読み込み
train, test = TabularDataset.splits(
path='./',
train='train.csv',
test='test.csv',
format='csv',
fields=[('en_text', EN_TEXT), ('ja_text', JA_TEXT)]
)
14. Vocabオブジェクト作成
これにより単語インデックス辞書/出現頻度/単語リスト、等が作成される。
JA_TEXT.build_vocab(train, test)
EN_TEXT.build_vocab(train, test)
14.1 日本語・英語双方の各単語にIDが付与されていることを確認
JA_TEXT.vocab.stoi
defaultdict(<function torchtext.vocab._default_unk_index>,
{'<unk>': 0,
'<pad>': 1,
'の': 2,
'は': 3,
'に': 4,
'て': 5,
'た': 6,
'を': 7,
'が': 8,
'.': 9,
'だ': 10,
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. 学習用バッチをイテレータオブジェクトとして作成
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 バッチの中身を確認
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は埋め込み層による分散表現の何番目の次元かを表している。
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モデル
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に限りなく近い出力とするためである。
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. モデルの初期化
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. 損失関数/最適化関数を定義
criterion = nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(), lr=0.0001)
21. 学習
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. 予測
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])