StatsBeginner: 初学者の統計学習ノート

統計学およびR、Pythonでのプログラミングの勉強の過程をメモっていくノート。たまにMacの話題。

小型のTransformerに英日翻訳をイチから学習させてみた

翻訳機を自作してみた

ChatGPTやDeepLの元になったTransformer*1をつかって、英語から日本語への翻訳を学習させてみました。
とりあえず現段階で、いくつかの英文を翻訳させてみた結果が以下のとおりです（青字が機械翻訳）。

Your time is limited, so don’t waste it living someone else’s life.
時間は限られているので、他人の人生を浪費しないでください。
（拙訳：君たちの時間は限られている。他人の人生を生きるようなことをしてそれを無駄にするな。／スティーブ・ジョブズ）
　
I have a dream that my four little children will one day live in a nation where they will not be judged by the color of their skin but by the content of their character.
私は4人の小さな子供たちが、いつか、自分の肌の色ではなく、その性格の特徴によって判断されない国家に住む夢を持っています。
（拙訳：私には夢がある。私の4人の子供たちがいつの日か、肌の色ではなく彼らの人格によって判断される国に暮らせるようになるという夢が。／キング牧師）
　
You are fake news!
偽のニュースです!
（拙訳：［CNNは］フェイクニュースだ！／トランプ大統領）
　
You may say I'm a dreamer. But I'm not the only one. I hope someday you'll join us. And the world will be as one.
夢想家だって言うかもしれないけど、私だけじゃない。いつか、ぼくらと合流してくれるといいんだけど。そして世界は、ひとつのものになる。
（拙訳：夢想家だと君は言うかもしれないけれど、僕は一人じゃない。いつか君も一緒になれれば。そして世界は一つになる。／ジョン・レノン）
　
The madman is not the man who has lost his reason. The madman is the man who has lost everything except his reason.
狂人は理性を失った人間ではない。狂人は、理性以外はすべてを失った男だ。
（安西徹雄訳：狂人とは理性を失った人のことではない。狂人とは、理性以外のあらゆる物を失った人のことである。／チェスタトン）
　
The safest general characterization of the European philosophical tradition is that it consists of a series of footnotes to Plato.
ヨーロッパ哲学伝統の最も安全な一般的な特徴は、それがプラトンへの一連の脚注から構成されていることです。
（拙訳：ヨーロッパ哲学の伝統について間違いなく言えるのは、その全てが、プラトン哲学へのひと続きの注釈に過ぎないということである。／ホワイトヘッド）
　
Violence sometimes may have cleared away obstructions quickly, but it never has proved itself creative.
暴力は時々すぐに妨害を片付けるかもしれませんが、それは創造的に証明されたことはありません。
（拙訳：暴力が、問題を手っ取り早く片付けるのに役立つことはある。しかし、暴力それ自身が創造的であったことは一度もない。／アインシュタイン）


表現がこなれてなかったり正しく訳せてなかったりするのですが、ひと昔前のGoogle翻訳ぐらいの性能は出ている気がしますね。個人で、小規模なコーパスを半日ぐらい学習させただけでも、これぐらい訳せるのだから、やはりTransformerは偉大ですね。ちなみに、手元のモデル自体もまだ改善の余地があり、コーパスも大きくしようと思えばできます。
　
　　

学習の概要

どういう学習をやったか簡単にメモしておくと、以下のような感じです。

  • 学習データについては、無償配布されている英日対訳のコーパスを何種類か集めて、クリーニングした上で統合したものを使いました（先日のエントリで手順を説明）。対訳文の数は、トータルで約128万ペアです。
  • モデルはAll You Need論文の基本構成に近いものですが、Transformerブロックはエンコーダ側・デコーダ側ともに3層ずつなので浅めです（All You Need論文は各6層）。フィードフォワード部分の隠れ層の次元も512と小さめです（All You Need論文は2048）。パラメータ数は約5990万個です（All You Need論文の構成だと約1.1億個になる）。
  • PyTorchのTransformerクラスを使っており、コードもPyTorchのチュートリアルに従いつつ、何点か修正を加えて動かしています。
  • Google Colab上で、A100というGPUのランタイムを使って学習させました。
  • 検証誤差はまだ低下中なので、まだ学習を進める余地があるのですが、とりあえず30エポックまで学習させて16時間ぐらいかかりました。


コードを書く上で、基本的にはPyTorchの以下のチュートリアルを参考にしましたが、次に述べるような変更を行っています。*2
Language Translation with nn.Transformer and torchtext — PyTorch Tutorials 2.0.1+cu117 documentation

  • コーパスは別のところから取ってきて、独自に整理したものを使った。
  • トークナイザはSentencePieceを使ってByte-pair Encodingを行って作成した。
  • 変数名がわかりにくかったり、記述の順番がわかりにくかったところは、適宜整理した（特にトークナイザ周辺）。
  • 学習時に勾配累積を行えるようにした。
  • 学習済みのモデルをロードして推論や追加学習を行うために、色々保存しておかないといけないものがあるので、整理して保存する部分を追加した。
  • 貪欲法だけでなくビームサーチでの推論もやってみた。


で、まだ改善の余地がありまして、具体的には以下のようなことが考えられます。

  • まず、30エポックで止めてますが、まだ検証誤差は減少中だったので、学習を継続すればもう少し精度が上がるかと思います。（⇒大して上がりませんでしたｗ）
  • コーパスは、全体としては3000万ぐらいのペアがある中から選別して128万件を抽出しているので、拡大することはできます。2500万と最大のデータ量をもつJParaCrawlというコーパスが、よくみると品質がいまいちなので、単純に全部使えばいいわけではなく精査が必要ですが*3、数百万ぐらいの規模にはできる気がします。
  • モデルをもっと大規模にする（トランスフォーマーブロックのスタックを増やす、フィードフォワード部分や埋め込みの次元数を増やす等）こともできます。コーパスを増やさずにモデル規模だけでどれだけ性能が上がるかは、やってみないとわかりませんが、AIl You Need論文のBase構成ぐらいまでは大きくしてもいいのかも知れない。
  • 単なる規模以外に、各種ハイパーパラメータの組み合わせも色々試して比較したわけではないので、工夫の余地があるかもです。
  • ビームサーチの精度が低いので、何か間違えてるかも知れません。冒頭で紹介した翻訳結果の例はすべて、貪欲法でやったものです。ビームサーチすると、近い内容ではあるものの、最後に。。。が繰り返されたりして変になります。
  • 大変そうなので今のところやる予定はないものの、たとえば日本語と英語のWikipediaをそれぞれLLMとして事前学習させた上で、ファインチューニングで翻訳を学習させると、精度は上がるんじゃないでしょうか。Wikipediaだったらデータを手に入れるのが簡単だし。ただし計算時間はえらいことになりそうで、私はできれば学習を自分を試したいのですが、Hugging Faceにある事前学習済みモデルを使えばいいかも知れません。
  • 翻訳の性能とは関係ないが、学習速度を早めるために、シーケンスをパディングするのをやめたほうがいいかもしれない。（今はバッチ内で最長のシーケンスに合わせてパディングしているので、バッチ内に長文が混ざると一気にテンソルがでかくなってメモリを圧迫している。）

　
　

学習に使ったコード

以下にコードを掲載しておきます。
モジュールをどこで使うのかがわかりやすいように、一番上にまとめてではなくセクションごとにインポートしてます。
セクション名の後に（☆）がついているのは、学習済みモデルを読み込んで推論を行う際に改めて実行する部分です。途中まで学習したモデルを読み込んで追加学習するときも、だいたい同じですが、当然学習に係る部分に要否は異なります。
いろいろ直しながら書いていったコードそのまんまなので、あまり整理されておらず、もっとラップしたりして見やすくしたほうがいいとは思います。
　

ランタイム立ち上げ、モジュールインポート（☆）

# SentencePieceのインストール
! pip install sentencepiece

# Googleドライブのマウント
from google.colab import drive
drive.mount('/content/drive')

# 作業ディレクトリ変更
import os
import sys
os.chdir('作業ディレクトリのパスを入れる')

# Colab側の保存フォルダをつくる
# Google Driveとの接続に不具合があったときを考えてColab側にもデータを一時保存するため
! mkdir /content/save_temp

　

ディレクトリの決定、パスの設定（☆）

これは完全に私の個人的な都合で設定しているものなので、どうでもいいっちゃどうでもいい。
コーパスとかモデルごとに「プロジェクト」としてまとめようと思ったので、そういうディレクトリ構成になっています。
学習済みもしくは途中まで学習したモデルを読み込んで、追加学習を行ったり推論を行ったりすることがあるわけですが、そのためには関連ファイルをきちんとロードする必要があって、最初は適当にやってたので何度かミスりました。それを避けるために、やや細かくディレクトリを整理しました。

import os

# プロジェクト名の設定
pj_name = 'EJTrans_1.3M'  # ★ここ設定

# プロジェクト自体のパス
pj_dir_path = 'projects/' + pj_name

# サブディレクトリのパス
corpus_dir_path   = pj_dir_path + '/corpus' # コーパス置き場
SP_dir_path       = pj_dir_path + '/SP'     # SentencePieceのモデル置き場
vocab_dir_path    = pj_dir_path + '/vocab'  # ボキャブラリの保存場所
arch_dir_path     = pj_dir_path + '/arch'   # モデルアーキテクチャの保存場所
model_dir_path    = pj_dir_path + '/model'  # 学習済みモデルのパラメータの保存場所

# ディレクトリの作成（なければ作る）
def check_and_make_dir(path):
  if not os.path.exists(path):
    os.makedirs(path)
check_and_make_dir(pj_dir_path)     # プロジェクト自体のディレクトリ
check_and_make_dir(corpus_dir_path) # コーパス置き場（英日別と、SP学習用の連結コーパス）
check_and_make_dir(SP_dir_path)     # SentencePieceのモデル置き場
check_and_make_dir(vocab_dir_path)  # ボキャブラリの保存場所
check_and_make_dir(arch_dir_path)   # モデルアーキテクチャの保存場所
check_and_make_dir(model_dir_path)  # 学習済みモデルの保存場所

# プロジェクトのログファイルを作成
import datetime
log_filename = f"{pj_dir_path}/training_log.txt"
current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
with open(log_filename, "w") as f:
    first_line = f"Training Log: created at {current_time}"
    f.write(first_line + "\n")

　

コーパスの事前処理

これもどういうコーパスを使うかによってかわりますが、私は先日のエントリで作成したような、「日本語と英語で1ファイルずつになっていて、行ごとに対訳として対応している」というコーパスファイルを作ってあったので、それを使う前提になっています。
ここでの処理としては、

  • 訓練データ、検証データ、テストデータに分割する（結果的に、テストは「適当に英文を入れてみて目視で性能をみる」ことしかやってないので、テストデータは使ってないですがｗ）
  • SentencePieceでバイトペアエンコーディングの学習をする際、All You Need論文にあわせて各言語共通の語彙空間をつくるようにしたので、英語と日本語を（SentencePieceの学習のためだけに）連結したものも作りました。
import shutil

'''リストの1要素を1行ずつテキストファイルに書き込む関数'''
def save_lines(text_list, path):
  with open(path, "w") as file:
      for l in text_list:
          _ = file.write(l)


'''オリジナルコーパスの取得元'''  # ★ここ設定
orig_corpus_en_path = 'data/en_ja_translation/mixed_1.3M_en.txt'
orig_corpus_ja_path = 'data/en_ja_translation/mixed_1.3M_ja.txt'


'''作成するコーパスのパスを予め指定'''
corpus_en_path = corpus_dir_path + '/corpus_en.txt'
corpus_ja_path = corpus_dir_path + '/corpus_ja.txt'


'''コーパスを元データ置き場から本プロジェクトのディレクトリにコピー配置'''
shutil.copyfile(orig_corpus_en_path, corpus_en_path)  # すでにあれば上書きする
shutil.copyfile(orig_corpus_ja_path, corpus_ja_path)  # すでにあれば上書きする


'''上記コーパスの読み込み'''
with open(corpus_en_path) as f:
    corpus_en = f.readlines()
with open(corpus_ja_path) as f:
    corpus_ja = f.readlines()


'''コーパスのサイズを限定する処理（デフォルトは全部利用）'''
full_corpus_len = len(corpus_en)
corpus_size = full_corpus_len  # サイズを限定する時は数値、全部使うならfull_corpus_len  ★ここ設定
if corpus_size < full_corpus_len:
  corpus_en = corpus_en[:corpus_size]
  corpus_ja = corpus_ja[:corpus_size]

# 後でトークナイザでロードするので（数を限定した場合は）改めて保存
if corpus_size < full_corpus_len:
  save_lines(corpus_en, corpus_en_path)
  save_lines(corpus_ja, corpus_ja_path)


'''コーパスを訓練用・検証用・テスト用に分割'''

# 分ける比率
split_ratio = [0.9, 0.05, 0.05]  # ★ここ設定

# 関数定義
def split_corpus(corpus, split_ratio, role, save_path):
  end_train = round(len(corpus)*split_ratio[0])
  end_valid = end_train + round(len(corpus)*split_ratio[1])
  if role=='train':
    save_lines(corpus[:end_train], save_path)
  if role=='valid':
    save_lines(corpus[end_train:end_valid], save_path)
  if role=='test':
    save_lines(corpus[end_valid:], save_path)

split_corpus(corpus_en, split_ratio, 'train', corpus_dir_path + '/corpus_en_train.txt')
split_corpus(corpus_en, split_ratio, 'valid', corpus_dir_path + '/corpus_en_valid.txt')
split_corpus(corpus_en, split_ratio, 'test',  corpus_dir_path + '/corpus_en_test.txt')
split_corpus(corpus_ja, split_ratio, 'train', corpus_dir_path + '/corpus_ja_train.txt')
split_corpus(corpus_ja, split_ratio, 'valid', corpus_dir_path + '/corpus_ja_valid.txt')
split_corpus(corpus_ja, split_ratio, 'test',  corpus_dir_path + '/corpus_ja_test.txt')

'''訓練コーパスの英語・日本語を連結したものを作成（SentencePiece用）'''
with open(corpus_dir_path + '/corpus_en_train.txt') as f:
    corpus_en_train = f.readlines()
with open(corpus_dir_path + '/corpus_ja_train.txt') as f:
    corpus_ja_train = f.readlines()
save_lines(corpus_en_train + corpus_ja_train, corpus_dir_path + '/corpus_enja_train.txt')

　

トークナイザの構築

トークナイザのセクションで必要な作業は、

  • BPEによって文章をサブワード（文字と単語の間ぐらいの中途半端な単位）に分割する
  • サブワードにIDをふって語彙空間を生成する
  • 開始記号、終了記号、未知語記号、パディング記号を定義して語彙空間にセットする

です。
なお、SentencePieceの学習にはけっこう時間がかかる（今回のデータ量で数十分）ので、1回構築したら、Transformerの学習中／済みモデルを読み込んで追加学習や推論を行う際は、SentencePieceのモデルも単に読み込むだけにしたほうがいいです。
チュートリアルの変数名は意味不明だったので変更しています。それ以外にもいろいろいじった気がする。

### SentencePieceの学習（コーパスがでかいとけっこう時間かかる）

import sentencepiece as spm

corppus_enja_path = corpus_dir_path + '/corpus_enja_train.txt'
sp_model_path = SP_dir_path + '/sp'
vocab_size = '30000'  # ★ここ設定

sp_command = '--input=' + corppus_enja_path + ' --model_prefix=' + sp_model_path + ' --vocab_size=' + vocab_size + ' --character_coverage=0.995 --model_type=bpe'
spm.SentencePieceTrainer.train(sp_command)


### トークナイザの構築

import sentencepiece as spm
import torch
from torchtext.vocab import build_vocab_from_iterator

'''SentencePieceの学習済みデータのロード'''
sp = spm.SentencePieceProcessor()
sp.load(SP_dir_path + '/sp.model')


'''src側とtgt側の言語の設定'''
SRC_LANGUAGE = 'en'
TGT_LANGUAGE = 'ja'


'''文字列をサブワードに分割する“text_to_subwords”の構築'''
# token_transformから名前を変更！！！

# サブワード分割器の生成器（所与のspモデルに基づく）
# こちらは与えられた「文字列」をサブワードに分けるもの
def sp_splitter(sp_model):
    def get_subwords(text):
        return sp_model.encode_as_pieces(text)
    return get_subwords

# 両言語のサブワード分割器を格納（今回は英日共通だが）し、言語名で取り出せるように
text_to_subwords = {}
text_to_subwords[SRC_LANGUAGE] = sp_splitter(sp)
text_to_subwords[TGT_LANGUAGE] = sp_splitter(sp)


'''サブワードをID化する“subwords_to_ids”の構築'''
# vocab_transformから名前を変更！！！

# SentencePieceの"encode_as_ids"でも作れるが、特殊記号のIDを明示的に
# 指定するために、build_vocab_from_iteratorを使っている

# ファイル内のテキストをサブワードに分割したものをイテレータとして返す関数
def file_to_subwords(data_path):
    with open(data_path, 'r', encoding='utf-8') as f:
        for line in f:
            yield sp.encode_as_pieces(line)

# 特殊トークンの定義（未知語、パディング、開始、終了）
UNK_IDX, PAD_IDX, BOS_IDX, EOS_IDX = 0, 1, 2, 3
special_symbols = ['<unk>', '<pad>', '<bos>', '<eos>']

# サブワードをIDに変換する変換器を両言語分格納
# これは語彙辞書としても機能する
subwords_to_ids = {}
subwords_to_ids[SRC_LANGUAGE] = build_vocab_from_iterator(file_to_subwords(corpus_dir_path + '/corpus_en_train.txt'), min_freq=1, specials=special_symbols, special_first=True)
subwords_to_ids[TGT_LANGUAGE] = build_vocab_from_iterator(file_to_subwords(corpus_dir_path + '/corpus_ja_train.txt'), min_freq=1, specials=special_symbols, special_first=True)

# デフォルトを未知語トークンに指定（この指定は必須）
for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:
  subwords_to_ids[ln].set_default_index(UNK_IDX)

'''ボキャブラリの保存'''
torch.save(subwords_to_ids[SRC_LANGUAGE], vocab_dir_path + '/vocab_en.pth')
torch.save(subwords_to_ids[TGT_LANGUAGE], vocab_dir_path + '/vocab_ja.pth')

　

トークナイザのロード（☆）

これは、学習中／済みモデルをロードして、学習を再開したり推論を行ったりする場合に使用する部分です。最初の学習時は実行しません（しても問題ないけど）。

# トークナイザのロード

import sentencepiece as spm
import torch

SRC_LANGUAGE = 'en'
TGT_LANGUAGE = 'ja'

'''サブワード分割器のロード'''
sp = spm.SentencePieceProcessor()
sp.load(SP_dir_path + '/sp.model')
def sp_splitter(sp_model):
    def get_subwords(text):
        return sp_model.encode_as_pieces(text)
    return get_subwords
text_to_subwords = {}
text_to_subwords[SRC_LANGUAGE] = sp_splitter(sp)
text_to_subwords[TGT_LANGUAGE] = sp_splitter(sp)

'''ID変換器（ボキャブラリ）のロード'''
subwords_to_ids = {}
subwords_to_ids[SRC_LANGUAGE] = torch.load(vocab_dir_path + '/vocab_en.pth')
subwords_to_ids[TGT_LANGUAGE] = torch.load(vocab_dir_path + '/vocab_ja.pth')

　

テキスト→テンソル変換器（☆）

テキストを、モデルに入れられる形のテンソルに変換するための部品です。チュートリアルではもっと下のほうに書いてあったのですが、分かりにくいので上にもってきています。

### テキストをTransformerに突っ込めるテンソルに変換する手続き
# チュートリアルではtext_transformと名付けられているもの

from torch.nn.utils.rnn import pad_sequence
from typing import List

# 与えられたトークンIDリストの先頭と末尾にそれぞれBOSとEOSのトークンIDを追加
BOS_IDX = subwords_to_ids[SRC_LANGUAGE]['<bos>']
EOS_IDX = subwords_to_ids[SRC_LANGUAGE]['<eos>']
def add_bos_eos(token_ids: List[int]):
    return torch.cat((torch.tensor([BOS_IDX]),
                      torch.tensor(token_ids),
                      torch.tensor([EOS_IDX])))

# テキストからテンソルまでの流れを定義
def sequential_transforms_src(input):
  tok_out = text_to_subwords[SRC_LANGUAGE](input)
  voc_out = subwords_to_ids[SRC_LANGUAGE](tok_out)
  ten_out = add_bos_eos(voc_out)
  return ten_out
def sequential_transforms_tgt(input):
  tok_out = text_to_subwords[TGT_LANGUAGE](input)
  voc_out = subwords_to_ids[TGT_LANGUAGE](tok_out)
  ten_out = add_bos_eos(voc_out)
  return ten_out

# テキストをtransformerに突っ込めるテンソルにする
text_to_tensor = {}
text_to_tensor[SRC_LANGUAGE] = sequential_transforms_src
text_to_tensor[TGT_LANGUAGE] = sequential_transforms_tgt


# 中身を一応確認する場合
text_to_tensor['en']('apple')

　

Transformerモデルの設計（☆）

　ここがモデルのコア部分です。

# Transformer構築に使うモジュール
from torch import Tensor
import torch
import torch.nn as nn
from torch.nn import Transformer
import math


'''ポジショナルエンコーディング（チュートリアル通り）'''

class PositionalEncoding(nn.Module):
    def __init__(self,
                 emb_size: int,
                 dropout: float,
                 maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        den = torch.exp(- torch.arange(0, emb_size, 2)* math.log(10000) / emb_size)
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        pos_embedding = torch.zeros((maxlen, emb_size))
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        pos_embedding = pos_embedding.unsqueeze(-2)

        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])


'''埋め込み（チュートリアル通り）'''

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size

    def forward(self, tokens: Tensor):
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)


'''Transformer本体（チュートリアル通り）'''

class Seq2SeqTransformer(nn.Module):
    def __init__(self,
                 num_encoder_layers: int,
                 num_decoder_layers: int,
                 emb_size: int,
                 nhead: int,
                 src_vocab_size: int,
                 tgt_vocab_size: int,
                 dim_feedforward: int = 512,
                 dropout: float = 0.1):
        super(Seq2SeqTransformer, self).__init__()
        self.transformer = Transformer(d_model=emb_size,
                                       nhead=nhead,
                                       num_encoder_layers=num_encoder_layers,
                                       num_decoder_layers=num_decoder_layers,
                                       dim_feedforward=dim_feedforward,
                                       dropout=dropout)
        self.generator = nn.Linear(emb_size, tgt_vocab_size)
        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
        self.positional_encoding = PositionalEncoding(
            emb_size, dropout=dropout)

    def forward(self,
                src: Tensor,
                trg: Tensor,
                src_mask: Tensor,
                tgt_mask: Tensor,
                src_padding_mask: Tensor,
                tgt_padding_mask: Tensor,
                memory_key_padding_mask: Tensor):
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None,
                                src_padding_mask, tgt_padding_mask, memory_key_padding_mask)
        return self.generator(outs)

    def encode(self, src: Tensor, src_mask: Tensor):
        return self.transformer.encoder(self.positional_encoding(
                            self.src_tok_emb(src)), src_mask)

    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        return self.transformer.decoder(self.positional_encoding(
                          self.tgt_tok_emb(tgt)), memory,
                          tgt_mask)

'''マスクの生成（チュートリアル通り）'''

PAD_IDX = subwords_to_ids[SRC_LANGUAGE]['<pad>']

'''
正方行列の上三角部分に True、それ以外に False を持つ行列を生成する。
その後、Trueの要素を-infに、Falseの要素を0.0に置き換える。
'''
def generate_square_subsequent_mask(sz):
    mask = (torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

def create_mask(src, tgt):
    src_seq_len = src.shape[0]
    tgt_seq_len = tgt.shape[0]

    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
    src_mask = torch.zeros((src_seq_len, src_seq_len),device=DEVICE).type(torch.bool)

    src_padding_mask = (src == PAD_IDX).transpose(0, 1)
    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
    return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

　

モデルのアーキテクチャ（Transformerのハイパーパラメータ）の設定

ここでアークテクチャと言っているのは、Transformerのモデルの各パーツのサイズみたいな感じです。
あとで呼び出すことがあるので、PyTorchのオブジェクトとして保存してます。

'''モデルアーキテクチャの設定と保存'''
# 学習時に使うドロップアウト率はclass Seq2SeqTransformerブロックの定義のところで固定してる

model_arch = {
    'NUM_ENCODER_LAYERS':3,                                # エンコーダ側のブロック数
    'NUM_DECODER_LAYERS':3,                                # デコーダ側のブロック数
    'EMB_SIZE':512,                                        # 埋め込み次元
    'NHEAD':8,                                             # マルチヘッドの数
    'SRC_VOCAB_SIZE':len(subwords_to_ids[SRC_LANGUAGE]),   # 英語の語彙次元
    'TGT_VOCAB_SIZE':len(subwords_to_ids[TGT_LANGUAGE]),   # 日本語の語彙次元
    'FFN_HID_DIM':512                                      # 順伝播層の次元
}
torch.save(model_arch, arch_dir_path + '/model_arch.pth')

　

学習の準備

Transformerモデルのアーキテクチャを設定してインスタンスを作成したり、オプティマイザのインスタンスを作成したりしています。学習率のスケジューリングを書くのが面倒だったので、固定にしてますが、エポックの進展にあわせて小さくしていくようにしたほうがいいとは思います。私は、途中まで学習して止めて、モデルをロードしてまた追加学習するということを繰り返していたので、手動で途中からlrを小さくしましたｗ

### デバイスの選択：GPUへの切り替え
# ifになってるが事実上、CPUで学習は不可能
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


'''(A) 最初から学習する場合の準備'''

# モデルのインスタンス構築
transformer = Seq2SeqTransformer(num_encoder_layers = model_arch['NUM_ENCODER_LAYERS'],
                                 num_decoder_layers = model_arch['NUM_DECODER_LAYERS'],
                                 emb_size = model_arch['EMB_SIZE'],
                                 nhead = model_arch['NHEAD'],
                                 src_vocab_size = model_arch['SRC_VOCAB_SIZE'],
                                 tgt_vocab_size = model_arch['TGT_VOCAB_SIZE'],
                                 dim_feedforward = model_arch['FFN_HID_DIM'])

# モデル内の重み・バイアスを一様分布で初期化
torch.manual_seed(0)
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)

# デバイスのセット
transformer = transformer.to(DEVICE)

# 誤差関数の定義（バッチ内で系列長をあわせるためのパディングトークンについては評価しない）
PAD_IDX = subwords_to_ids[SRC_LANGUAGE]['<pad>']
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# ネットワークの初期パラメータを渡し、学習率等も設定してオプティマイザを初期化  ★ここ設定する
optimizer = torch.optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)  # epsはゼロで割るのを回避するための部品


'''(B) 途中から学習を再開する場合の準備'''

# アーキテクチャ設定のロード
model_arch = torch.load(arch_dir_path + '/model_arch.pth')

# インスタンス立ち上げ
transformer = Seq2SeqTransformer(num_encoder_layers = model_arch['NUM_ENCODER_LAYERS'],
                                 num_decoder_layers = model_arch['NUM_DECODER_LAYERS'],
                                 emb_size = model_arch['EMB_SIZE'],
                                 nhead = model_arch['NHEAD'],
                                 src_vocab_size = model_arch['SRC_VOCAB_SIZE'],
                                 tgt_vocab_size = model_arch['TGT_VOCAB_SIZE'],
                                 dim_feedforward = model_arch['FFN_HID_DIM'])

# モデルにデバイスをセット
transformer.to(DEVICE)

# 学習済みパラメータをロード
transformer.load_state_dict(torch.load(model_dir_path + '/best_model_epoch=25_valloss=2.400.pth'))

# 誤差関数の定義（バッチ内で系列長をあわせるためのパディングトークンについては評価しない）
PAD_IDX = subwords_to_ids[SRC_LANGUAGE]['<pad>']
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# ネットワークの初期パラメータを渡し、学習率等も設定してオプティマイザを初期化  ★ここ設定する　（学習再開時にlrを下げたりしてる）
optimizer = torch.optim.Adam(transformer.parameters(), lr=0.00001, betas=(0.9, 0.98), eps=1e-9)  # epsはゼロで割るのを回避するための部品

　

データローダのセット

学習時に、コーパスの中からデータを取って、バッチにまとめてモデルに供給していく装置です。

### データローダの構築
# 学習ルーチンにデータを供給するローダ

from torch.utils.data import Dataset
from torch.utils.data import DataLoader

# バッチサイズ  ★ここ設定する
BATCH_SIZE = 128

# コーパスファイルを読み込んで1件ずつ与えていく動きの関数
class CustomDataset(Dataset):
    def __init__(self, src_path, tgt_path):
        with open(src_path, 'r', encoding='utf-8') as f:
            self.src_data = [line.strip() for line in f]
        with open(tgt_path, 'r', encoding='utf-8') as f:
            self.tgt_data = [line.strip() for line in f]

    def __len__(self):
        return len(self.src_data)

    def __getitem__(self, idx):
        return self.src_data[idx], self.tgt_data[idx]

# テキストをテンソルに変換しパディングする関数
PAD_IDX = subwords_to_ids[SRC_LANGUAGE]['<pad>']

def collate_fn(batch):
    src_batch, tgt_batch = [], []
    for src_sample, tgt_sample in batch:
        src_batch.append(text_to_tensor[SRC_LANGUAGE](src_sample.rstrip("\n")))
        tgt_batch.append(text_to_tensor[TGT_LANGUAGE](tgt_sample.rstrip("\n")))
    src_batch = pad_sequence(src_batch, padding_value=PAD_IDX)
    tgt_batch = pad_sequence(tgt_batch, padding_value=PAD_IDX)
    return src_batch, tgt_batch

# 訓練用のローダを構築
train_dataset = CustomDataset(corpus_dir_path + '/corpus_en_train.txt', corpus_dir_path + '/corpus_ja_train.txt')
train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)

# 検証用のローダを構築
val_dataset = CustomDataset(corpus_dir_path + '/corpus_en_valid.txt', corpus_dir_path + '/corpus_ja_valid.txt')
val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)

　

学習ルーチン

学習と検証のルーチンです。勾配累積を使うようにチュートリアルから変更してます。

### 学習ルーチン

from torch.utils.data import DataLoader

# 勾配累積の回数を設定
GRADIENT_ACCUMULATION_STEPS = 2

'''新しい訓練ルーチン'''
# 勾配累積を行うように変更した

def train_epoch(model, optimizer, train_dataloader):
    model.train()
    losses = 0

    for step, (src, tgt) in enumerate(train_dataloader):
        src = src.to(DEVICE)
        tgt = tgt.to(DEVICE)

        tgt_input = tgt[:-1, :]

        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

        logits = model(src, tgt_input, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask, src_padding_mask)

        tgt_out = tgt[1:, :]
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
        loss = loss / GRADIENT_ACCUMULATION_STEPS  # 勾配累積のためのスケーリング
        loss.backward()
        losses += loss.item()

        # 指定したステップ数で勾配を累積した後、パラメータを更新
        if (step + 1) % GRADIENT_ACCUMULATION_STEPS == 0:
            optimizer.step()
            optimizer.zero_grad()

    return losses / len(list(train_dataloader))


'''検証ルーチン'''
def evaluate(model, val_dataloader):
    model.eval()
    losses = 0

    for src, tgt in val_dataloader:
        src = src.to(DEVICE)
        tgt = tgt.to(DEVICE)

        tgt_input = tgt[:-1, :]

        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

        logits = model(src, tgt_input, src_mask, tgt_mask,src_padding_mask, tgt_padding_mask, src_padding_mask)

        tgt_out = tgt[1:, :]
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
        losses += loss.item()

    return losses / len(list(val_dataloader))

　

学習の実施

以下で学習を進めるのですが、要するに
train_loss = train_epoch(transformer, optimizer, train_dataloader)
の部分でローダからデータがモデルに供給されて、勾配が計算されたりパラメータが更新されたりしていきます。
検証ロスがベストを更新したら、モデルを保存するようにしています。その際、Google Colabの不具合でGoogle Driveとの接続がおかしくなることがあるので、Colabサーバ側でもモデルを保存するようにしています。これは緊急時にここからモデルを救うというものなので、別に保存してなくてもいいといえばいいです。
あと、ログを書き込むようにしてます。

'''学習の実施'''
from timeit import default_timer as timer

NUM_EPOCHS = 35  # ★ここ設定する

# 初期の最低検証ロスを無限大として設定
best_val_loss = float('inf')
#best_val_loss = 2.355  # 学習再開時は　★ここ設定する

for epoch in range(1, NUM_EPOCHS+1):  # 学習再開時でエポックのカウントを続きからにしたい場合は　★ここ設定する　（通常は1から）
    torch.cuda.empty_cache()
    start_time = timer()
    train_loss = train_epoch(transformer, optimizer, train_dataloader)
    end_time = timer()
    val_loss = evaluate(transformer, val_dataloader)
    log_message = (f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Val loss: {val_loss:.3f}, "f"Epoch time = {(end_time - start_time):.3f}s")
    print(log_message)
    with open(log_filename, "a") as f:
      f.write(log_message + "\n")

    # 検証ロスが最低を更新したらモデルを保存
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(transformer.state_dict(), "/content/save_temp/best_model.pth")  # Colabにも一時保存（上書き）
        torch.save(transformer.state_dict(), f"{model_dir_path}/best_model_epoch={epoch}_valloss={val_loss:.3f}.pth")  # Google Driveにモデル保存

　

保存したモデルのロード（☆）

ここは、追加学習や推論時に何が必要かをちゃんと理解して、整理しておく必要があります。

'''①以下のセクションをすべて実行する'''
# ランタイム立ち上げ、モジュールインポート（☆）
# ディレクトリの決定、パスの設定（☆）
# トークナイザのロード（☆）
# テキスト→テンソル変換器（☆）
# Transformerモデルの設計（☆）

'''②モデルアーキテクチャ（ハイパーパラメータ）のロード'''
model_arch = torch.load(arch_dir_path + '/model_arch.pth')

'''③モデルのインスタンス構築'''
transformer = Seq2SeqTransformer(num_encoder_layers = model_arch['NUM_ENCODER_LAYERS'],
                                 num_decoder_layers = model_arch['NUM_DECODER_LAYERS'],
                                 emb_size = model_arch['EMB_SIZE'],
                                 nhead = model_arch['NHEAD'],
                                 src_vocab_size = model_arch['SRC_VOCAB_SIZE'],
                                 tgt_vocab_size = model_arch['TGT_VOCAB_SIZE'],
                                 dim_feedforward = model_arch['FFN_HID_DIM'])
transformer.to(DEVICE)

'''④保存済みパラメータをロード'''
transformer.load_state_dict(torch.load(model_dir_path + '/best_model.pth'))

　

パラメータ数とトークン数の確認

別にどうでもいいと言えばどうでもいいですが、モデルのパラメータ数やデータのトークン数をカウントしたいとき用。

# モデルのパラメータ数確認
num_params = sum(p.numel() for p in transformer.parameters())
print(num_params)

# コーパスのトークン数確認（訓練データのみ）
with open(corpus_dir_path + '/corpus_en_train.txt') as f:
    corpus_en = f.readlines()
with open(corpus_dir_path + '/corpus_ja_train.txt') as f:
    corpus_ja = f.readlines()

num_token_en = 0
for text in corpus_en:
    token_ids = subwords_to_ids[SRC_LANGUAGE](text_to_subwords[SRC_LANGUAGE](text))
    num_token_en += len(token_ids)
num_token_ja = 0
for text in corpus_ja:
    token_ids = subwords_to_ids[TGT_LANGUAGE](text_to_subwords[TGT_LANGUAGE](text))
    num_token_ja += len(token_ids)


print("Number of English tokens:", num_token_en)
print("Number of English tokens:", num_token_ja)
print("Number of pairs of parallel corpus:", len(corpus_en))

　

推論の関数（貪欲法）

推論を実施するための関数です。

'''貪欲法での推論の関数'''

# 出力のトークン列を貪欲法で求める
def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(DEVICE)
    src_mask = src_mask.to(DEVICE)

    memory = model.encode(src, src_mask)
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(DEVICE)
    for i in range(max_len-1):
        memory = memory.to(DEVICE)
        tgt_mask = (generate_square_subsequent_mask(ys.size(0))
                    .type(torch.bool)).to(DEVICE)
        out = model.decode(ys, memory, tgt_mask)
        out = out.transpose(0, 1)
        prob = model.generator(out[:, -1])
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.item()

        ys = torch.cat([ys,
                        torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
        if next_word == EOS_IDX:
            break
    return ys

# 上記関数をつかって翻訳文を生成するラッパー
def translate(model: torch.nn.Module, src_sentence: str):
    model.eval()
    src = text_to_tensor[SRC_LANGUAGE](src_sentence).view(-1, 1)
    num_tokens = src.shape[0]
    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
    tgt_tokens = greedy_decode(
        model,  src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()
    return " ".join(subwords_to_ids[TGT_LANGUAGE].lookup_tokens(list(tgt_tokens.cpu().numpy()))).replace("<bos>", "").replace("<eos>", "").replace(" ", "").replace("▁", "")


'''ビームサーチでの推論の関数'''

import heapq

# ビームサーチにより出力のトークン列を求める
def beam_search_decode(model, src, src_mask, max_len, start_symbol, beam_size):
    src = src.to(DEVICE)
    src_mask = src_mask.to(DEVICE)

    memory = model.encode(src, src_mask)
    initial_input = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(DEVICE)

    # 初期仮説を設定
    hypotheses = [(0.0, initial_input)]
    completed_hypotheses = []

    for _ in range(max_len - 1):
        new_hypotheses = []

        for score, prev_output in hypotheses:
            prev_output = prev_output.to(DEVICE)
            tgt_mask = (generate_square_subsequent_mask(prev_output.size(0))
                        .type(torch.bool)).to(DEVICE)

            out = model.decode(prev_output, memory, tgt_mask)
            out = out.transpose(0, 1)
            prob = model.generator(out[:, -1])

            topk_probs, topk_indices = torch.topk(prob, beam_size)

            for prob, next_word in zip(topk_probs[0], topk_indices[0]):
                new_score = score + prob.item()
                new_output = torch.cat([prev_output, next_word.view(1, 1)], dim=0)
                new_hypotheses.append((new_score, new_output))

        # ビーム内の仮説をスコアでソートして上位のものを保持
        new_hypotheses.sort(reverse=True, key=lambda x: x[0])
        hypotheses = new_hypotheses[:beam_size]

        # EOSトークンがあるかチェック
        completed_hypotheses = [hypothesis for hypothesis in hypotheses if hypothesis[1][-1].item() == EOS_IDX]
        if len(completed_hypotheses) >= beam_size:
            break

    # 最もスコアの高い仮説を取得
    best_hypothesis = max(hypotheses, key=lambda x: x[0])[1]

    return best_hypothesis.flatten()

# 上記関数をつかって翻訳文を生成するラッパー
def translate_beam_search(model: torch.nn.Module, src_sentence: str, beam_size: int):
    model.eval()
    src = text_to_tensor[SRC_LANGUAGE](src_sentence).view(-1, 1)
    num_tokens = src.shape[0]
    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
    tgt_tokens = beam_search_decode(
        model, src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX, beam_size=beam_size).flatten()
    return " ".join(subwords_to_ids[TGT_LANGUAGE].lookup_tokens(list(tgt_tokens.cpu().numpy()))).replace("<bos>", "").replace("<eos>", "").replace(" ", "").replace("▁", "")

　

推論（翻訳）の実施

いよいよ翻訳をやってみます。

### 推論の実施
# 以下の例文はコーパスに含まれていないことを確認済み。


eng = ['I am a pen.',
       'Your time is limited, so don’t waste it living someone else’s life.',
       'I have a dream that my four little children will one day live in a nation where they will not be judged by the color of their skin but by the content of their character.',
       'You are fake news!',
       'You may say I\'m a dreamer. But I\'m not the only one. I hope someday you\'ll join us. And the world will be as one.',
       'The madman is not the man who has lost his reason. The madman is the man who has lost everything except his reason.',
       'We are being afflicted with a new disease of which some readers may not yet have heard the name, but of which they will hear a great deal in the years to come—namely, technological unemployment. This means unemployment due to our discovery of means of economising the use of labour outrunning the pace at which we can find new uses for labour.  But this is only a temporary phase of maladjustment. All this means in the long run that mankind is solving its economic problem. I would predict that the standard of life in progressive countries one hundred years hence will be between four and eight times as high as it is. There would be nothing surprising in this even in the light of our present knowledge. It would not be foolish to contemplate the possibility of afar greater progress still.'
]

# 貪欲法
for e in eng:
  print(e + '\n' + translate(transformer, e) + '\n')

# ビームサーチ
for e in eng:
  print(e + '\n' + translate_beam_search(transformer, e, 5) + '\n')

*1:有名なのでここで説明する必要はないと思うが、深層学習（ディープラーニング）のアーキテクチャで、2017年にGoogleの“Attention is all you need”という論文で提案され、自然言語処理の性能を革命的に向上させた。

*2:いわゆるLLMの事前学習を行うためのチュートリアルもあって、こちらでは、エンコーダ側とデコーダ側が分かれたクラスが使われてる。まだ詳しくは見てない。Language Modeling with nn.Transformer and torchtext — PyTorch Tutorials 2.0.1+cu117 documentation

*3:以下のブログ記事ではJParaCrawlを全部使って学習した例が紹介されているが、大して性能がよくはなっていない模様。SageMakerでJParaCrawlのコーパスを使って翻訳モデルを作成する