PyTorchを使ってSeq2seqの実装をしてみたので、その実装方法について書いてみます。学習に用いるデータセットには実装したSeq2seqが正しく学習できているか簡単に評価できる、トイ・プロブレム（足し算）を使っています。
Google Colaboratoryで動作確認をしたので、誰でも簡単にプログラムを実行できるようにしています。
環境
今回書いたプログラムはGoogle Colaboratoryで実装から実行まですることができるので環境構築に手間取ることはないと思います。さらにGoogle ColaboratoryのGPUを利用するので短時間で学習からテストをし、実装したSeq2seqの性能を見ることができます。
Google Colaboratoryの導入の仕方についてはこちら。
Seq2seq
プログラムを実装する前にSeq2seqについて少し書いてみます。
Seq2seqとは Sequence to Sequence モデル、モデルの入力と出力が時系列データになっており、時系列データを別の時系列データに変換してくれます。
今回は最大で３桁の数字どうしを足し算する式を入力とし、出力にはその計算結果を与えます。例えば、入力は１００＋１０のような式の場合、出力にはその計算結果である１１０を与えるようにします。Seq2seqモデルは１００＋１０が与えられたとき計算結果の１１０が出力されるように学習していきます。
Seq2seqはEncoder Decoder モデルとも呼ばれ、イメージとしては以下のようになります。
Encoderに入力として変換したいデータを渡しEncoderでそれを処理したのちDecoderに処理結果を渡して入力したデータの変換結果をDecoderが出力します。
EncoderとDecoderにはRNNが使われており与えられた時系列データをそれぞれ処理します。
上の例ではRNNとしてGRUを使っています。GRUではなく、シンプルなRNNやLSTMを用いることも可能です。
Encoderには入力として文字単位か単語単位でデータを与えますが、今回は与える入力データを文字単位に分割してEncoderに与えます。入力データが１００＋１０だとすると、”１”、”０”、”０”、”＋”、”１”、”０”に分割します。
ただ、ここで注意しなければならないことがあります。
今回は最大で3桁の数字の足し算を入力とします。つまり、入力データには７３＋２や１１２＋９９９が与えられます。ここで問題となるのが入力データの長さの異なる可変長のデータという点です。７３＋２や１１２＋９９９を文字単位に分割すると長さは４と７になります。
今回実装するプログラムではミニバッチ学習をするので入力データの形状を揃える必要があります。そのための方法としてパディングを使います。パディングをすることにより、可変長のデータの長さを全て同じ長さに揃えミニバッチ学習が可能になります。
入力データはパディングを使って長さを７に統一します。長さが７未満の場合、余ったところにはパディングを表す記号 <pad> を入れることにします。
７３＋２にパディングを適用すると、
”７”、”３”、”＋”、”２”、”<pad>”、”<pad>”、”<pad>”
１１２＋９９９だと、
”１”、”１”、”２”、”＋”、”９”、”９”、”９”
のようになります。
Encoderに入力データを与えると、ｈ（隠れ状態）が出力されます。この隠れ状態には入力したデータから正しい出力結果を得るための必要な情報に変換したものが入っており、これを初期状態としてDecoderに渡しています。
DecoderはEncoderから渡された情報と教師データから損失を計算していき学習を進めていきます。
学習データの用意
学習データを用意します。今回は２万のデータを作成し訓練用とテスト用に分けて使用します。学習データを用意するコードは下に置いてあります。ここでは、コードの説明をざっくりしていきます。
まず、文字がkey、文字IDがvalueの辞書を用意します。
word2id = {str(i): i for i in range(10)}
word2id.update({"<pad>": 10, "+": 11, "<eos>": 12})
この辞書には０～９の数字とパディングのための<pad>、演算子＋、区切り文字である<eos>が登録されています。（※区切り文字は後ほど出てきます。）
Encoderに入力データを与えるときは１００＋１０のような式を文字単位で分解し、辞書をもとにしてIDに変換します。
生成したデータセットのひとつを見てみると、
print(train_x[0])
このデータは入力データで３＋７を変換したものになります。
print(train_t[0])
このデータは出力データで計算結果の１０を変換したものになります。１２は区切り文字の<eos>です。
EncoderとDecoderの実装
PyTorchを使ってEncoderとDecoderの実装をしていきます。コード全体は下に置いてあります。
Encoderのイメージとしてはこんな感じです。
３＋１１を入力としたときのEncoderです。文字IDに変換したものをEmbedding Layerに渡していきます。
このイメージをコードで実装すると以下になります。
class Encoder(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, batch_size=100):
super(Encoder, self).__init__()
self.hidden_dim = hidden_dim
self.batch_size = batch_size
self.word_embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=word2id["<pad>"])
self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)
def forward(self, indices):
embedding = self.word_embeddings(indices)
if embedding.dim() == 2:
embedding = torch.unsqueeze(embedding, 1)
_, state = self.gru(embedding, torch.zeros(1, self.batch_size, self.hidden_dim, device=device))
return state
Encoderでは受け取った文字IDをEmbeding Layerに与え単語ベクトルに変換したものをGRUに渡しています。GRUの出力statestateを出力します。このstateには入力データを変換するための必要な情報が入っており、これをDecoderへと渡します。
DecoderはEncoderから出力されたstateと教師データをもとに目的とする出力を生成します。
出力が１０のときのDecoderの図です。１２は<eos>、区切り文字です。
この図をコードで実装すると以下のようになります。
class Decoder(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, batch_size=100):
super(Decoder, self).__init__()
self.hidden_dim = hidden_dim
self.batch_size = batch_size
self.word_embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=word2id["<pad>"])
self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)
self.output = nn.Linear(hidden_dim, vocab_size)
def forward(self, index, state):
embedding = self.word_embeddings(index)
if embedding.dim() == 2:
embedding = torch.unsqueeze(embedding, 1)
gruout, state = self.gru(embedding, state)
output = self.output(gruout)
return output, state
区切り文字を入れるには２つの理由があります。
- Decoderに文字列の生成開始を知らせる。
- 文字列生成の終了を知らせる。
出力を１０としたとき、文字IDに変換すると[12, 1, 0, 12, 10, 10] になります。今回は３桁どうしの足し算をしているので計算結果は最大で４桁になり、ここに区切り文字を入れることにより出力の長さは６になります。
出力が１０のときのDecoderには教師データとして[12, 1, 0, 12, 10, 10] を使いますが、入力には[12, 1, 0, 12, 10] を与え、それに対応する出力として[1, 0, 12, 10, 10] を与えます。
学習
コードは下に置いてあります。だいだい１００epochでそれなりの正答率がとれます。
学習コードはここのPyTorchのTutorialを参考にして実装しました。
Encoderには入力データを文字列単位で与えていますが、Decoderには１文字ずつ与えています。
テスト
コードは下に置いてあります。１００epochで正答率はだいだい８割くらいです。１０epochごとに保存したモデルにテストデータを渡してテストをしています。
学習するときのパラメータをいじったり、epoch数を増やせばもっといい結果が得られるかもしれません。
データの用意
import random
from sklearn.model_selection import train_test_split
word2id = {str(i): i for i in range(10)}
word2id.update({"<pad>": 10, "+": 11, "<eos>": 12})
id2word = {v: k for k, v in word2id.items()}
def load_dataset(N=20000):
def generate_number():
number = [random.choice(list("0123456789")) for _ in range(random.randint(1, 3))]
return int("".join(number))
def padding(string, training=True):
string = "{:*<7s}".format(string) if training else "{:*<6s}".format(string)
return string.replace("*", "<pad>")
def transform(string, seq_len=7):
tmp = []
for i, c in enumerate(string):
try:
tmp.append(word2id[c])
except:
tmp += [word2id["<pad>"]] * (seq_len - i)
break
return tmp
data = []
target = []
for _ in range(N):
x = generate_number()
y = generate_number()
z = x + y
left = padding(str(x) + "+" + str(y))
right = padding(str(z), training=False)
data.append(transform(left))
right = transform(right, seq_len=6)
right = [12] + right[:5]
right[right.index(10)] = 12
target.append(right)
return data, target
data, target = load_dataset()
train_x, test_x, train_t, test_t = train_test_split(data, target, test_size=0.1)
EncoderとDecoder
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
embedding_dim = 16
hidden_dim = 128
vocab_size = len(word2id)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class Encoder(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, batch_size=100):
super(Encoder, self).__init__()
self.hidden_dim = hidden_dim
self.batch_size = batch_size
self.word_embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=word2id["<pad>"])
self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)
def forward(self, indices):
embedding = self.word_embeddings(indices)
if embedding.dim() == 2:
embedding = torch.unsqueeze(embedding, 1)
_, state = self.gru(embedding, torch.zeros(1, self.batch_size, self.hidden_dim, device=device))
return state
class Decoder(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, batch_size=100):
super(Decoder, self).__init__()
self.hidden_dim = hidden_dim
self.batch_size = batch_size
self.word_embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=word2id["<pad>"])
self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)
self.output = nn.Linear(hidden_dim, vocab_size)
def forward(self, index, state):
embedding = self.word_embeddings(index)
if embedding.dim() == 2:
embedding = torch.unsqueeze(embedding, 1)
gruout, state = self.gru(embedding, state)
output = self.output(gruout)
return output, state
encoder = Encoder(vocab_size, embedding_dim, hidden_dim).to(device)
decoder = Decoder(vocab_size, embedding_dim, hidden_dim).to(device)
criterion = nn.CrossEntropyLoss(ignore_index=word2id["<pad>"])
encoder_optimizer = optim.Adam(encoder.parameters(), lr=0.001)
decoder_optimizer = optim.Adam(decoder.parameters(), lr=0.001)
学習する
from datetime import datetime
from sklearn.utils import shuffle
batch_size=100
def train2batch(data, target, batch_size=100):
input_batch = []
output_batch = []
data, target = shuffle(data, target)
for i in range(0, len(data), batch_size):
input_tmp = []
output_tmp = []
for j in range(i, i+batch_size):
input_tmp.append(data[j])
output_tmp.append(target[j])
input_batch.append(input_tmp)
output_batch.append(output_tmp)
return input_batch, output_batch
def get_current_time():
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
print("Training...")
n_epoch = 100
for epoch in range(1, n_epoch+1):
input_batch, output_batch = train2batch(train_x, train_t)
for i in range(len(input_batch)):
encoder_optimizer.zero_grad()
decoder_optimizer.zero_grad()
inputs = torch.tensor(input_batch[i], device=device)
outputs = torch.tensor(output_batch[i], device=device)
encoder_hidden = encoder(inputs)
source = outputs[:, :-1]
target = outputs[:, 1:]
decoder_hidden = encoder_hidden
loss = 0
for i in range(source.size(1)):
decoder_output, decoder_hidden = decoder(source[:, i], decoder_hidden)
decoder_output = torch.squeeze(decoder_output)
loss += criterion(decoder_output, target[:, i])
loss.backward()
encoder_optimizer.step()
decoder_optimizer.step()
if epoch % 10 == 0:
print(get_current_time(), "Epoch %d: %.2f" % (epoch, loss.item()))
if epoch % 10 == 0:
model_name = "seq2seq_calculator_v{}.pt".format(epoch)
torch.save({
'encoder_model': encoder.state_dict(),
'decoder_model': decoder.state_dict(),
}, model_name)
print("Saving the checkpoint...")
テストする
import numpy as np
result = """---------------
Q:{:>10s}
A:{:>10s}
T/F: {}
---------------"""
encoder = Encoder(vocab_size, embedding_dim, hidden_dim, batch_size=1).to(device)
decoder = Decoder(vocab_size, embedding_dim, hidden_dim, batch_size=1).to(device)
for epoch in range(10, 101, 10):
model_name = "seq2seq_calculator_v{}.pt".format(epoch)
checkpoint = torch.load(model_name)
encoder.load_state_dict(checkpoint["encoder_model"])
decoder.load_state_dict(checkpoint["decoder_model"])
print("Checkpoint {:>3d}".format(epoch))
print("-"*30)
accuracy = 0
with torch.no_grad():
for i in range(len(test_x)):
x = test_x[i]
input_tensor = torch.tensor([x], device=device)
state = encoder(input_tensor)
token = "<eos>"
try:
padded_idx_x = x.index(word2id["<pad>"])
except ValueError:
padded_idx_x = len(x)
left = "".join(map(lambda c: str(id2word[c]), x[:padded_idx_x]))
right = []
for _ in range(7):
index = word2id[token]
input_tensor = torch.tensor([index], device=device)
output, state = decoder(input_tensor, state)
prob = F.softmax(torch.squeeze(output))
index = torch.argmax(prob.cpu().detach()).item()
token = id2word[index]
if token == "<eos>":
break
right.append(token)
right = "".join(right)
flag = ["F", "T"][eval(left) == int(right)]
if flag == "T":
accuracy += 1
print("Accuracy: {:.2f}".format(accuracy / len(test_x)))
print("-"*30)
Deep Learningによる自然言語処理を勉強したい方へのおすすめの本です。
tensorflow, chainerやPyTorchといったフレームワークを使わずにゼロからnumpyを使ってディープラーニングの実装をしています。
扱っている内容はword2vec, RNN, GRU, seq2seqやAttentionなど、、、
Google Colaboratoryでプログラムを動かす
最後にGoogle Colaboratoryを使った学習方法を紹介します。（※ランタイムのタイプはPython 3、GPUに設定しておきます。）
まずは、drive内のディレクトリをマウントします。
www.pytry3g.com
以下のコードをセルにコピペして実行します。
from google.colab import drive
drive.mount('/content/drive')
これによりdriveのMy Drive以下にGoogle Driveのルートディレクトリがマウントされます。
cd drive/My\ Drive/任意の作業ディレクトリ
次にGoogle Driveにある任意の作業ディレクトリに移動します。この作業ディレクトリには学習が進んでいくとモデルが保存されます。
!pip install torch > /dev/null
次にPyTorchをインストールします。
あとはソースコードを順に実行していくだけです。
