transformersのTokenizerで固定長化する - Money Forward Developers Blog

灼熱からこんにちは、CTO室AI推進部の@ken11です。PC排熱にあぶられ干からびています。自宅のPCで学習をするということ、それは自宅がホットアイルになるということ…←

最近弊社のエンジニアブログの投稿数が多くて嬉しい限りです。きっと僕がこうやってしょうもないことばかり書いてるからエンジニアブログを書くことへのハードルが下がっているのだと信じています← 嘘です、それは技術広報のluccaさんのおかげです。

ところでみんな、自然言語処理やってる？日本語の問題は日本の人しか解決しないので我々が頑張るしかないんだけど、大丈夫そ？冗談はさておき、そういう思いもあって自分は自然言語処理が好きです。日頃から公私ともにいろいろチャレンジしているのですが、今日はその自然言語処理における入力の話です。失敗談です。。

自然言語処理の学習における入力の固定長化

BERTでもなんでも、なにかしら深層学習系の自然言語処理モデルの学習をやってみたことがある人はだいたいわかると思いますが、入力は固定長化することが多いです。 (逆に自分は可変長で問題ないケースを知らない)

text = ["こんにちはリヴァイアさん", "こんにちは"]
tokenized_text = tokenizer(text)  # なにかしらのtokenize処理
print(tokenized_text)
# [[3, 4, 5], [3]]

このように長さが異なると困るので

LENGTH = 12
padding_text = padding(tokenized_text, length=LENGTH)  # なにかしらのpadding処理
print(padding_text)
# [[3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0], [3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]

なにかしらのpadding処理を施すことが一般的かと思います。この例では長さが 12 になるように 0 で埋めていますね。例えばよく使われるBERTのモデルだと、デフォルト長が512になってるので無意識に512にpaddingしていたりするのではないでしょうか。

paddingする際に、なにで埋めるのか(例では0でしたが、これは使用するボキャブラリーなどシーンによって異なる)、右詰or左詰等、考慮すべき事項は多いものだったりもします。また、同様にtruncationもあります。よく例に挙がるBERTのデフォルトが512と(それなりに)長いこともあって、このtruncationはあまり例が少ないかもしれないですね。

transformersのTokenizerでpadding/truncation

さて、このようなpadding/truncation処理はkerasのpad_sequence等で行うこともあると思いますが、自然言語処理ライブラリとして有名なhuggingfaceのtransformersでもその処理が用意されています。

PreTrainedTokenizerBaseの呼び出し時にパラメータ指定できるようになっており、各Tokenizerはこのクラスを継承しているはずなので、transformersで各Tokenizerを利用している場合はこのパラメータを指定することでpadding/truncationができます。

パラメータ指定時の挙動

とはいえ実はここに癖があって、自分はハマりました。 ……まあ、上記ドキュメント読めばその通りなんですが。

まずpaddingの方ですが、Trueにしたとてmax_lengthでpaddingされるわけじゃないです。

LENGTH = 12
tokenizer = なにかtransformersのTokenizer
text = ["こんにちはリヴァイアさん", "こんにちは"]
tokenized_text = tokenizer(text, return_tensors="pt", max_length=LENGTH, padding=True)

たとえばこのようなコードを実行したとき、普通に考えると tokenized_text['input_ids'] は

[[3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0], [3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]

のようになっていることを期待してしまいます。 max_length指定してますし、paddingもTrueですし。しかしこの結果は以下のようになります。

print(tokenized_text['input_ids'])
# [[3, 4, 5], [3, 0, 0]]

padding=True の挙動は padding="longest" となっており、これは要素の中で最も長いものにあわせてpaddingするというものです。そう、この場合 max_length を指定していても意味が無いのです…

期待通り長さを 12 にするには以下のように設定します。

LENGTH = 12
tokenizer = なにかtransformersのTokenizer
text = ["こんにちはリヴァイアさん", "こんにちは"]
tokenized_text = tokenizer(text, return_tensors="pt", max_length=LENGTH, padding="max_length")

このように padding="max_length" を指定することで、

print(tokenized_text['input_ids'])
# [[3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0], [3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]

期待された長さのものを得ることができます。

では、truncationの方はどうでしょうか？例として、input_idの長さが20くらいになる "なにかとてつもなく長い文章" を追加してみます。仮にinput_idが

[100, 103, 389, 222, 43, 89, 4, 788, 100, 103, 389, 222, 43, 89, 4, 788, 43, 89, 4, 788]

になるとしましょう。

tokenizer = なにかtransformersのTokenizer
text = ["こんにちはリヴァイアさん", "こんにちは", "なにかとてつもなく長い文章"]
tokenized_text = tokenizer(text, return_tensors="pt")
print(tokenized_text['input_ids'])
# [[3, 4, 5], [3], [100, 103, 389, 222, 43, 89, 4, 788, 100, 103, 389, 222, 43, 89, 4, 788, 43, 89, 4, 788]]

この入力でtruncationを使ってみます

LENGTH = 12
tokenizer = なにかtransformersのTokenizer
text = ["こんにちはリヴァイアさん", "こんにちは", "なにかとてつもなく長い文章"]
tokenized_text = tokenizer(text, return_tensors="pt", max_length=LENGTH, truncation=True)
print(tokenized_text['input_ids'])
# [[3, 4, 5], [3], [100, 103, 389, 222, 43, 89, 4, 788, 100, 103, 389, 222]]

短くなってますね。そうです、truncationはTrueを指定すればmax_lengthの長さにカットされます。 paddingの挙動をふまえると "shortest" 的なものかと思っていたらそうではなかった。

つまり、いろいろな長さの文章の入力を固定長化したい場合はこれらのパラメータを以下のように組み合わせるとよさそうです。

LENGTH = 12
tokenizer = なにかtransformersのTokenizer
text = ["こんにちはリヴァイアさん", "こんにちは", "なにかとてつもなく長い文章"]
tokenized_text = tokenizer(text, return_tensors="pt", max_length=LENGTH, truncation=True, padding="max_length")
print(tokenized_text['input_ids'])
# [[3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0], [3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [100, 103, 389, 222, 43, 89, 4, 788, 100, 103, 389, 222]]

max_length=LENGTH, truncation=True, padding="max_length" とすることで無事にすべての要素を 12 でそろえることができました。

失敗例

では実際にpaddingとtruncationのパラメータ指定を間違えるとどうなるのか見ていきましょう。

LENGTH = 12
tokenizer = なにかtransformersのTokenizer
text = ["こんにちはリヴァイアさん", "こんにちは", "なにかとてつもなく長い文章"]
tokenized_text = tokenizer(text, return_tensors="pt", max_length=LENGTH, padding=True)

これは最初の僕の失敗。 max_length指定しておけばtruncateされると思っていた(←)し、max_lengthにpaddingもされると思っていた… 結果はこうなってしまいます

print(tokenized_text['input_ids'])
# [[3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [100, 103, 389, 222, 43, 89, 4, 788, 100, 103, 389, 222, 43, 89, 4, 788, 43, 89, 4, 788]]

truncateはされないしその結果最も長いものが20になるのでpaddingは全部20にあわせて行われてしまいます。これ今は例なので20程度ですけど、実際にほんとにとてつもなく長い文章が紛れ込んでいると、そのときだけメモリ使用量がスパイクしてあっけなくOOM死します。 max_length を指定して満足していましたがこのパターンではmax_lengthはなんの意味もないしメモリ使用量が読めずに頭を抱えるだけです。ひどいミスですね…ドキュメントはちゃんと読みましょう(自戒)

では続いてこちらではどうでしょうか？

LENGTH = 12
tokenizer = なにかtransformersのTokenizer
text = ["こんにちはリヴァイアさん", "こんにちは", "なにかとてつもなく長い文章"]
tokenized_text = tokenizer(text, return_tensors="pt", max_length=LENGTH, padding=True, truncation=True)

よさそうですね？ちゃんとmax_lengthにtruncateされますし。 …しかしどうでしょう？試行ごとに必ずtruncateされる長さの文章が存在する状況であれば、実質すべてmax_lengthになるのでいいかもしれません。しかし、truncateされる長さの文章がたまにしか現れない場合、試行によって最大長のバラツキが出ます。それで問題が起こることは少ないかもしれませんが、たとえば以下のようなケースだと困るシーンが出てくると思います。

# max_length = 12とする
# 1step目の入力
[[3, 4, 5], [3, 0, 0], [8, 8, 0]]
# 2step目の入力
[[6, 6, 5], [3, 2, 0], [22, 3, 4]]
# 3step目の入力
[[6, 6, 5, 8, 12], [5, 0, 0, 0, 0], [9, 7, 0, 0, 0]]
# 4step目の入力
[[3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0], [3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [100, 103, 389, 222, 43, 89, 4, 788, 100, 103, 389, 222]]

このように、基本的に短い文章でたまに長い文章が流れてくるものだと、試行ごとにサイズにバラツキが発生しており、やはりメモリ使用量を見誤りやすいのではないでしょうか。きちんとすべてがmax_lengthで流れてくる前提で見積もればいいですが、最初の数回の試行だけ見て見積もってしまうとOOM死します。

このように幾度かの失敗を乗り越えてたどり着いたのが max_length=LENGTH, truncation=True, padding="max_length" でした。