torchtextの仕様変更対応 (2) TabularDataset BucketIterator

PyTorch 1.11(torchtext 0.12)より自然言語処理で活用していたtorchtextのFieldやTabularDatasetなど便利な機能がなくなりました。PyTorch 1.10(torchtext 0.11)まではlegacyに移動されいましたが利用することは可能でした。しかし、PyTorch 1.11(torchtext 0.12)で完全に削除されてしまいました。

今回は、TabularDataset,BucketIteratorについて確認していきたいと思います。

基本的な使い方

テキストファイルの読み込みからミニバッチ学習までサポートしてくれる便利な機能でした。
以下のファイルを例に基本的な利用方法を見ていきます。
(textとlabelの間はタブです。)

test.txt

text	label
あ	1
あ い	2
あ い う	3
あ い え	3
あ い う え お	5
あ い う え お か	6
あ い う え お き	6
あ い う え お か き く	8
あ い う え お か き く け	9

旧実装

データ読み込み

TabularDatasetを利用すれば、表形式のファイルを読み込むことができます。
自然言語処理の場合は、文中に,が含まれるため、タブ区切りのtsvが使われることが多いと思います。
まず、各項目の処理を設定するため各項目ごとにFieldを設定します。一つ目は、テキストの格納されたFieldを、二つ目は、ラベル用のFieldを用意しておきます。Fieldは、既定値では、スペースで区切るtokenizerが設定されています。
TabularDatasetでファイルを読み込みデータセットにします。pathにファイルのパス、formatは、tsvに設定します。filedsは先ほど定義したTEXT,LABELを設定します。

from torchtext.legacy.data import Field, LabelField, TabularDataset, BucketIterator

# Firld定義
TEXT = Field(batch_first=True)
LABEL = LabelField()
# データ読み込み
ds = TabularDataset(path="test.txt", format='tsv', fields=[('text', TEXT), ('label', LABEL)], skip_header=True)

データセットの中身を確認してみましょう。テキストは、スペースごとに区切られてリストに格納されています。

for i in range(len(ds)):
    print(ds[i].text, ds[i].label)

['あ'] 1
['あ', 'い'] 2
['あ', 'い', 'う'] 3
['あ', 'い', 'え'] 3
['あ', 'い', 'う', 'え', 'お'] 5
['あ', 'い', 'う', 'え', 'お', 'か'] 6
['あ', 'い', 'う', 'え', 'お', 'き'] 6
['あ', 'い', 'う', 'え', 'お', 'か', 'き', 'く'] 8
['あ', 'い', 'う', 'え', 'お', 'か', 'き', 'く', 'け'] 9

単語辞書作成

単語辞書を作成し、辞書の内容を確認します。

# 単語辞書作成
TEXT.build_vocab(ds)

TEXT.vocab.stoi

            {'<unk>': 0,
             '<pad>': 1,
             'あ': 2,
             'い': 3,
             'う': 4,
             'え': 5,
             'お': 6,
             'か': 7,
             'き': 8,
             'く': 9,
             'け': 10})

ラベルを同様に作成し、各ラベルの出現回数を表示します。

LABEL.build_vocab(ds)

LABEL.vocab.freqs

Counter({'1': 1, '2': 1, '3': 2, '5': 1, '6': 2, '8': 1, '9': 1})

ミニバッチ学習

ミニバッチ学習用にイテレータを生成します。DataLoaderに相当します。
エポックごとにシャッフルするように、shuffle=Trueを設定しています。(train=Trueと設定することも可能)

# イテレータ生成
biter = BucketIterator(dataset=ds, shuffle=True, batch_size=3)

各バッチごとのデータを確認してみます。

for i, (texts, labels) in enumerate(biter):
    print(i)
    for text, label in zip(texts, labels):
        print(text, LABEL.vocab.itos[label])

0
tensor([2, 3, 4, 5, 6, 7]) 6
tensor([2, 3, 1, 1, 1, 1]) 2
tensor([2, 1, 1, 1, 1, 1]) 1
1
tensor([2, 3, 4, 5, 6, 7, 8, 9]) 8
tensor([2, 3, 4, 5, 6, 8, 1, 1]) 6
tensor([2, 3, 4, 5, 6, 1, 1, 1]) 5
2
tensor([2, 3, 4, 1, 1, 1, 1, 1, 1]) 3
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 5, 1, 1, 1, 1, 1, 1]) 3

データは、辞書を利用し文字から数値に変換されています。
ミニバッチごとに、データ長が統一されていることがわかります。不足部分は1(<pad>)でパディングされています。
もう２回実行してみます。

0
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 4, 5, 6, 8, 1, 1, 1]) 6
tensor([2, 3, 4, 5, 6, 7, 8, 9, 1]) 8
1
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 5, 1, 1, 1]) 3
tensor([2, 3, 4, 5, 6, 7]) 6
2
tensor([2, 3, 4]) 3
tensor([2, 1, 1]) 1
tensor([2, 3, 1]) 2

0
tensor([2, 3, 5, 1, 1, 1, 1, 1]) 3
tensor([2, 3, 4, 5, 6, 7, 1, 1]) 6
tensor([2, 3, 4, 5, 6, 7, 8, 9]) 8
1
tensor([2, 1, 1, 1, 1, 1, 1, 1, 1]) 1
tensor([2, 3, 4, 1, 1, 1, 1, 1, 1]) 3
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
2
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 4, 5, 6, 8]) 6
tensor([2, 3, 1, 1, 1, 1]) 2

エポックごとにシャッフルされていることがわかります。

まとめ

これだけで、ファイルの読み込みから辞書作成、ミニバッチ学習用のデータ生成まで行ってくれます。非常に便利でした。

from torchtext.legacy.data import Field, LabelField, TabularDataset, BucketIterator

# Firld定義
TEXT = Field(batch_first=True)
LABEL = LabelField()
# データ読み込み
ds = TabularDataset(path="test.txt", format='tsv', fields=[('text', TEXT), ('label', LABEL)], skip_header=True)
# 単語辞書作成
TEXT.build_vocab(ds)
LABEL.build_vocab(ds)
# イテレータ生成
biter = BucketIterator(dataset=ds, shuffle=True, batch_size=3)

新実装

TabularDataset,BucketIteratorを使わずに実装していきます。

データ読み込み

TabularDatasetがなくなり、ファイルの読み込みを自分で実装する必要があります。
タブ区切りのファイルの読み込みには、いろんな方法がありますが、pandasを利用しました。read_tableを利用すれば、タブ区切りのファイルも簡単に読み込むことができます。

import pandas as pd

# データ読み込み
df = pd.read_table('test.txt')

print(df)

                text  label
0                  あ      1
1                あ い      2
2              あ い う      3
3              あ い え      3
4          あ い う え お      5
5        あ い う え お か      6
6        あ い う え お き      6
7    あ い う え お か き く      8
8  あ い う え お か き く け      9

pandasのデータフレームにデータが格納されました。

単語分割

単語に分割するために、tokenizerを作成します。今回は、スペースで区切るだけのためget_tokenizerを利用しました。tokenizerにNoneを設定すれば単にスペースで区切る処理を行ってくれます。

from torchtext.data.utils import get_tokenizer

# tokenizer設定
tokenizer = get_tokenizer(tokenizer=None)

テキストをスペースで区切ります。

# 単語分割
df['text'] = df['text'].map(lambda x: tokenizer(x))

print(df)

                          text  label
0                          [あ]      1
1                       [あ, い]      2
2                    [あ, い, う]      3
3                    [あ, い, え]      3
4              [あ, い, う, え, お]      5
5           [あ, い, う, え, お, か]      6
6           [あ, い, う, え, お, き]      6
7     [あ, い, う, え, お, か, き, く]      8
8  [あ, い, う, え, お, か, き, く, け]      9

スペースで区切られました。

単語辞書作成

単語辞書を作成します。前回説明したCounterを用いてもよいのですが、Counterから辞書作成まで一気に行ってくれるbuild_vocab_from_iteratorを利用しました。

from torchtext.vocab import build_vocab_from_iterator

# 単語辞書作成
text_vocab = build_vocab_from_iterator(df['text'], specials=('<unk>', '<pad>'))
text_vocab.set_default_index(text_vocab['<unk>'])

辞書の内容を確認します。

text_vocab.get_stoi()

{'け': 10,
 'い': 3,
 '<unk>': 0,
 'う': 4,
 '<pad>': 1,
 'あ': 2,
 'え': 5,
 'お': 6,
 'か': 7,
 'き': 8,
 'く': 9}

ラベルも同様に行います。
build_vocab_from_iteratorは、文字型を想定しているようです。今回は、ラベルに数値で文字数を入れていたため文字型に変換しています。もともとクラス分類などでクラス名を文字型で設定している場合は変換は不要です。

# ラベル辞書
df['label'] = df['label'].astype(str)  # ラベルが数値のため文字型に変換。文字型で指定していた場合はこの処理は不要
label_vocab = build_vocab_from_iterator(df['label'])

label_vocab.get_stoi()

{'9': 6, '5': 4, '3': 0, '8': 5, '6': 1, '1': 2, '2': 3}

transform設定

テキスト、ラベルの変換方法を定義します。以前は、Fieldを定義しておけば、BucketIteratorで行ってくれていましたが、自前で変換する必要があります。
テキストは、辞書による変換とパディング、Tensor型への変換を行います。パディングは、ミニバッチごとに系列長を統一するため不足部分がパディングされます。
ラベルは、辞書による変換とTensor型への変換を行います。

import torchtext.transforms as T

# transform生成
text_transform = T.Sequential(
    T.VocabTransform(text_vocab),
    T.ToTensor(padding_value=text_vocab['<pad>'])
)
label_transform = T.Sequential(
    T.VocabTransform(label_vocab),
    T.ToTensor()
)

ミニバッチ学習

ミニバッチ学習用にDataLoaderに作成します。BucketIteratorでは、内部で変換を行ってくれていましたが、DataLoaderでは明示的に変換を指示する必要があります。
変換を行う関数では、ミニバッチのデータを受け取りテキスト、ラベルそれぞれtransformにて変換を行います。

# ミニバッチ時のデータ変換関数
def collate_batch(batch):
    texts = text_transform([text for (text, label) in batch])
    labels = label_transform([label for (text, label) in batch])
    return texts, labels

エポックごとにシャッフルするようにshuffle=True、データ変換を行うためcollate_fnに先ほど定義した関数を設定しています。

from torch.utils.data import DataLoader

# DataLoader設定
data_loader = DataLoader(df.values, batch_size=3, shuffle=True, collate_fn=collate_batch)

各バッチごとのデータを確認してみます。

for i, (texts, labels) in enumerate(data_loader):
    print(i)
    for text, label in zip(texts, labels):
        print(text, label_vocab.lookup_token(label))

0
tensor([2, 1, 1, 1, 1, 1, 1, 1, 1]) 1
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 5, 1, 1, 1, 1, 1, 1]) 3
1
tensor([2, 3, 4, 5, 6, 8, 1, 1]) 6
tensor([2, 3, 4, 5, 6, 7, 8, 9]) 8
tensor([2, 3, 1, 1, 1, 1, 1, 1]) 2
2
tensor([2, 3, 4, 1, 1, 1]) 3
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 4, 5, 6, 7]) 6

ミニバッチごとに、旧実装と同様にデータ長が統一されていることがわかります。不足部分はtransformで設定した1(<pad>)でパディングされています。
もう２回実行してみます。

0
tensor([2, 3, 4, 5, 6, 8]) 6
tensor([2, 3, 4, 5, 6, 7]) 6
tensor([2, 3, 1, 1, 1, 1]) 2
1
tensor([2, 3, 4, 1, 1, 1, 1, 1, 1]) 3
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 4, 5, 6, 7, 8, 9, 1]) 8
2
tensor([2, 3, 5, 1, 1]) 3
tensor([2, 1, 1, 1, 1]) 1
tensor([2, 3, 4, 5, 6]) 5

0
tensor([2, 3, 4, 5, 6, 8, 1, 1, 1]) 6
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 4, 5, 6, 1, 1, 1, 1]) 5
1
tensor([2, 3, 5, 1, 1, 1, 1, 1]) 3
tensor([2, 3, 4, 5, 6, 7, 8, 9]) 8
tensor([2, 3, 4, 1, 1, 1, 1, 1]) 3
2
tensor([2, 1, 1, 1, 1, 1]) 1
tensor([2, 3, 1, 1, 1, 1]) 2
tensor([2, 3, 4, 5, 6, 7]) 6

旧実装と同様の結果になりました。

まとめ

torchtextの提供する機能はできるだけ利用したつもりですが、結構複雑ですね。
今回は、データの読み込みにpandasを利用しました。また、データ変換にtransformを利用しました。もっと簡単に実装する方法があればぜひ教えてください。

import pandas as pd
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
import torchtext.transforms as T
from torch.utils.data import DataLoader

# データ読み込み
df = pd.read_table('test.txt')

# tokenizer設定
tokenizer = get_tokenizer(tokenizer=None)
# 単語分割
df['text'] = df['text'].map(lambda x: tokenizer(x))
# 単語辞書作成
text_vocab = build_vocab_from_iterator(df['text'], specials=('<unk>', '<pad>'))
text_vocab.set_default_index(text_vocab['<unk>'])
# ラベル辞書
df['label'] = df['label'].astype(str)  # ラベルが数値のため文字型に変換。文字型で指定していた場合はこの処理は不要
label_vocab = build_vocab_from_iterator(df['label'])
# transform生成
text_transform = T.Sequential(
    T.VocabTransform(text_vocab),
    T.ToTensor(padding_value=text_vocab['<pad>'])
)
label_transform = T.Sequential(
    T.VocabTransform(label_vocab),
    T.ToTensor()
)
# ミニバッチ時のデータ変換関数
def collate_batch(batch):
    texts = text_transform([text for (text, label) in batch])
    labels = label_transform([label for (text, label) in batch])
    return texts, labels
# DataLoader設定
data_loader = DataLoader(df.values, batch_size=3, shuffle=True, collate_fn=collate_batch)

単語数によるソート

自然言語処理の学習時、できるだけ単語数が同じ文章をまとめて学習すると良いと言われています。長さが異なると最長の単語数に合わせてパディングが行われます。単語数がほぼ同じであればパディングを抑えることができます。

旧実装

BucketIteratorでsort_within_batch=Trueとすれば、単語数ごとに並び替えてミニバッチを作成してくれます。
BucketIteratorを以下のように変更し試してみます。

# イテレータ生成
biter = BucketIterator(dataset=ds, shuffle=True, sort_within_batch=True, sort_key=lambda x: len(x.text), batch_size=3)

各バッチごとのデータを確認してみます。

for i, (texts, labels) in enumerate(biter):
    print(i)
    for text, label in zip(texts, labels):
        print(text, LABEL.vocab.itos[label])

単語数が近いものごとにまとめられました。

0
tensor([2, 3, 4, 5, 6, 8]) 6
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 5, 1, 1, 1]) 3
1
tensor([2, 3, 4]) 3
tensor([2, 3, 1]) 2
tensor([2, 1, 1]) 1
2
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 4, 5, 6, 7, 8, 9, 1]) 8
tensor([2, 3, 4, 5, 6, 7, 1, 1, 1]) 6

もう２回実行してみます。

0
tensor([2, 3, 4, 5, 6, 8]) 6
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 5, 1, 1, 1]) 3
1
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 4, 5, 6, 7, 8, 9, 1]) 8
tensor([2, 3, 4, 5, 6, 7, 1, 1, 1]) 6
2
tensor([2, 3, 4]) 3
tensor([2, 3, 1]) 2
tensor([2, 1, 1]) 1

0
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 4, 5, 6, 7, 8, 9, 1]) 8
tensor([2, 3, 4, 5, 6, 7, 1, 1, 1]) 6
1
tensor([2, 3, 4]) 3
tensor([2, 3, 1]) 2
tensor([2, 1, 1]) 1
2
tensor([2, 3, 4, 5, 6, 8]) 6
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 5, 1, 1, 1]) 3

非常に簡単に実装できました。

新実装

新実装では、自前でミニバッチごとに返却するデータを決める必要があります。
Samplerでミニバッチごとに返却するデータのインデックスを決めます。Samplerには、ミニバッチごとにインデックスを返却する__iter__と長さを求める__len__を実装します。
__iter__では、単語数ごとにソートしミニバッチごとのインデックスを決めています。また、返却する順番をシャッフルします。__len__は、データ数をバッチサイズで割り切り上げます。

class SortSampler():
    def __init__(self, df, batch_size, shuffle=True):
        self.df = df
        self.batch_size = batch_size
        self.shuffle = shuffle

    def __iter__(self):
        import random
        # 文章中の単語数でソート可能なように単語数を求める。
        indices = [(i, len(text)) for i, text in enumerate(self.df['text'])]
        # 同じ単語数の文章がシャッフルされるようにシャッフルを行う。
        random.shuffle(indices)
        # 単語数でソート
        pooled_indices = sorted(indices, key=lambda x: x[1])
        # ソートに用いた単語数を削除
        pooled_indices = [x[0] for x in pooled_indices]
        # バッチサイズごとに分割
        mini_batch = []
        for i in range(0, len(pooled_indices), self.batch_size):
            mini_batch.append(pooled_indices[i:i + self.batch_size])
        # ミニバッチごとにシャッフル
        if self.shuffle:
            random.shuffle(mini_batch)
        # ミニバッチのindexを返却
        for i in range(len(mini_batch)):
            yield mini_batch[i]

    def __len__(self):
        from math import ceil
        return ceil(len(df) / self.batch_size)

サンプラーを生成します。

sort_sampler = SortSampler(df, batch_size=3)

データローダーを生成します。ミニバッチへの分割やシャッフルはサンプラーで行うため、batch_sizeやshuffleは指定しません。

data_loader = DataLoader(df.values, collate_fn=collate_batch, batch_sampler=sort_sampler)

実行してみます。

for i, (texts, labels) in enumerate(data_loader):
    print(i)
    for text, label in zip(texts, labels):
        print(text, label_vocab.lookup_token(label))

0
tensor([2, 1, 1]) 1
tensor([2, 3, 1]) 2
tensor([2, 3, 5]) 3
1
tensor([2, 3, 4, 1, 1, 1]) 3
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 4, 5, 6, 7]) 6
2
tensor([2, 3, 4, 5, 6, 8, 1, 1, 1]) 6
tensor([2, 3, 4, 5, 6, 7, 8, 9, 1]) 8
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9

0
tensor([2, 3, 4, 5, 6, 8, 1, 1, 1]) 6
tensor([2, 3, 4, 5, 6, 7, 8, 9, 1]) 8
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
1
tensor([2, 1, 1]) 1
tensor([2, 3, 1]) 2
tensor([2, 3, 5]) 3
2
tensor([2, 3, 4, 1, 1, 1]) 3
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 4, 5, 6, 7]) 6

0
tensor([2, 3, 5, 1, 1, 1]) 3
tensor([2, 3, 4, 5, 6, 1]) 5
tensor([2, 3, 4, 5, 6, 8]) 6
1
tensor([2, 3, 4, 5, 6, 7, 1, 1, 1]) 6
tensor([2, 3, 4, 5, 6, 7, 8, 9, 1]) 8
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
2
tensor([2, 1, 1]) 1
tensor([2, 3, 1]) 2
tensor([2, 3, 4]) 3

旧実装と同様の結果になりました。

しかし、かなり複雑です。もっと簡単に実装できるか知りたいです。
PyTorch側でのサポートを拡大してもらいたいところです。