torchtextの仕様変更対応 (3) TorchData利用

PyTorch 1.11(torchtext 0.12)より自然言語処理で活用していたtorchtextのFieldやTabularDatasetなど便利な機能がなくなりました。前回の「torchtextの仕様変更対応 (2) TabularDataset BucketIterator」では、pandasを利用し実装しました。
今回は、PyTorch 1.11と同時にリリースされたTorchDataを利用し実装してみます。TorchData(0.3.0)は現時点ではβ版です。

データ

前回と同様のデータを使って試します。
textとlabelの間はタブです。

text.txt

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

旧実装

参考のため旧実装です。詳細は、前回を記事（「torchtextの仕様変更対応 (2) TabularDataset BucketIterator」）を参照してください。

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)

新実装

TorchDataを利用し実装していきます。
今後は、PyTorchのデータ操作関連、DataLoaderなどはTorchDataが主流になるようです。
TorchDataは、CSVやjsonの読み込みやWebからの読み込み、zipファイルの展開など様々な機能をサポートしておりデータパイプとして接続していくことで利用します。
詳細は、TorchDataのドキュメントを参照してください。

データ読み込み

データパイプを接続しながら読み込みを行います。
ファイルをオープンするFileOpener、CSV(TSV)を読み込むCSVParserを利用します。
FileOpenerの第一引数は、datapipeを指定する必要があるため、単一ファイルパスを渡すことができないようです。リストとしてファイルパスを渡します。
CSVParserでCSVを読み込みます。ここでは、タブ区切りのためdelimiterを指定します。ヘッダを読み飛ばすため、skip_lines=1を指定します。csv.readerのパラメータが利用できます。
datapipeを順番につないでいけばよいです。

import torchdata.datapipes as dp

datapipe = dp.iter.FileOpener(['test.txt'], mode='rt')
datapipe = dp.iter.CSVParser(datapipe, delimiter='\t', skip_lines=1)

データが読み込めたか確認します。

for text, lable in datapipe:
    print(text, lable)

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

問題なく読み込めていることがわかります。

単語分割

次に、単語分割を行います。ここでは、単にスペースで区切るだけのためget_tokernizerを用いました。
Mapperで変換を行います。変換を行う関数を定義します。ここではlambda式で関数を定義しています。テキスト部分の変換を行うためinput_colに0を設定しています。
ここでもdatapipeをつなげるだけです。

from torchtext.data.utils import get_tokenizer

# tokenizer設定
tokenizer = get_tokenizer(tokenizer=None)
# 単語分割
datapipe = dp.iter.Mapper(datapipe, lambda text: tokenizer(text), input_col=0)

実行後にlambda式を利用しているためワーニングが表示されます。TorchDataのドキュメントにも記載されているので無視して構わないでしょう。
データを確認します。

for text, lable in datapipe:
    print(text, lable)

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

問題なく単語ごとに分割できました。

単語辞書作成

単語辞書は、文章の方を利用し作成します。１つ目の項目ですが、datapipeから１つ目の項目を取り出す必要がありますが、適切なデータパイプが見つかりませんでした。Mapperを利用し１つ目だけ取り出します。

datapipe_text = dp.iter.Mapper(datapipe, fn=lambda x: x[0])

データを確認してみます。

for text in datapipe_text:
    print(text)

['あ']
['あ', 'い']
['あ', 'い', 'う']
['あ', 'い', 'え']
['あ', 'い', 'う', 'え', 'お']
['あ', 'い', 'う', 'え', 'お', 'か']
['あ', 'い', 'う', 'え', 'お', 'き']
['あ', 'い', 'う', 'え', 'お', 'か', 'き', 'く']
['あ', 'い', 'う', 'え', 'お', 'か', 'き', 'く', 'け']

ちゃんとテキストの部分だけ取り出せています。

前回と同様にbuild_vocab_from_iteratorを利用します。datapipeを渡せば大丈夫です。

from torchtext.vocab import build_vocab_from_iterator

# 単語辞書作成
text_vocab = build_vocab_from_iterator(datapipe_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}

ラベルも同様に行います。datapipeを利用し２つめの項目を取り出し辞書を作成します。

datapipe_label = dp.iter.Mapper(datapipe, fn=lambda x: x[1])
# ラベル辞書
label_vocab = build_vocab_from_iterator(datapipe_label)

label_vocab.get_stoi()

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

transform設定

テキスト、ラベルの変換方法を定義します。この部分は前回と同様ですので、詳細は前回の記事を参照してください。

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()
)

ミニバッチ学習

まずは、前回同様にミニバッチごとに辞書を用いてデータを変換する関数を定義します。

# ミニバッチ時のデータ変換関数
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にdatapipeを渡せばよいのですが、datapipeでは順番にデータを読み込むだけでシャッフルが行えません。そこで、to_map_style_datasetを用いてmapに変換します。これでいったんデータをすべて読み込むのでシャッフルに対応できます。

from torch.utils.data import DataLoader
from torchtext.data.functional import to_map_style_dataset

# mapに変換
ds = to_map_style_dataset(datapipe)
# DataLoader設定
data_loader = DataLoader(ds, 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, 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, 1, 1]) 6
tensor([2, 1, 1, 1, 1, 1, 1, 1]) 1
tensor([2, 3, 4, 5, 6, 7, 8, 9]) 8
2
tensor([ 2,  3,  4,  5,  6,  7,  8,  9, 10]) 9
tensor([2, 3, 1, 1, 1, 1, 1, 1, 1]) 2
tensor([2, 3, 4, 1, 1, 1, 1, 1, 1]) 3

もう２回実行した結果です。

0
tensor([2, 3, 5, 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
1
tensor([2, 3, 4, 1, 1]) 3
tensor([2, 3, 4, 5, 6]) 5
tensor([2, 3, 1, 1, 1]) 2
2
tensor([2, 3, 4, 5, 6, 7]) 6
tensor([2, 1, 1, 1, 1, 1]) 1
tensor([2, 3, 4, 5, 6, 8]) 6

0
tensor([2, 3, 4, 5, 6, 8]) 6
tensor([2, 3, 4, 5, 6, 7]) 6
tensor([2, 1, 1, 1, 1, 1]) 1
1
tensor([2, 3, 1, 1, 1, 1, 1, 1]) 2
tensor([2, 3, 4, 5, 6, 7, 8, 9]) 8
tensor([2, 3, 5, 1, 1, 1, 1, 1]) 3
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, 4, 5, 6, 1, 1, 1, 1]) 5

問題ないですね。

まとめ

datapipeは、クラスを利用し定義してきましたが、関数を利用する方が推奨されています。それぞれ関数がありますのでドキュメントを参照してください。

datapipe = dp.iter.FileOpener(['test.txt'], mode='rt')
datapipe = dp.iter.CSVParser(datapipe, delimiter='\t', skip_lines=1)
datapipe = dp.iter.Mapper(datapipe, lambda text: tokenizer(text), input_col=0)

変更後です。

datapipe = dp.iter.FileOpener(['test.txt'], mode='rt').parse_csv(delimiter='\t', skip_lines=1).map(lambda text: tokenizer(text), input_col=0)

すべてのプログラムです。

import torchdata.datapipes as dp
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
from torchtext.data.functional import to_map_style_dataset

# tokenizer設定
tokenizer = get_tokenizer(tokenizer=None)
# データ読み込み
datapipe = dp.iter.FileOpener(['test.txt'], mode='rt').parse_csv(delimiter='\t', skip_lines=1).map(lambda text: tokenizer(text), input_col=0)
# 単語辞書作成
datapipe_text = datapipe.map(fn=lambda x: x[0])
text_vocab = build_vocab_from_iterator(datapipe_text, specials=('<unk>', '<pad>'))
text_vocab.set_default_index(text_vocab['<unk>'])
# ラベル辞書
datapipe_label = datapipe.map(fn=lambda x: x[1])
label_vocab = build_vocab_from_iterator(datapipe_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
# mapに変換
ds = to_map_style_dataset(datapipe)
# DataLoader設定
data_loader = DataLoader(ds, batch_size=3, shuffle=True, collate_fn=collate_batch)

参考

to_map_style_datasetで一旦、mapに変換しましたが、shuffleを利用すればdatapipeのままで対応できそうです。内部的に読み込んでシャッフルしているようです。
こちらを利用してもよいかもしれません。

# シャッフル対応
s_datapipe = datapipe.shuffle()
# DataLoader設定
data_loader = DataLoader(s_datapipe, batch_size=3, shuffle=True, collate_fn=collate_batch)

今後、datapipeが主流になりそうなので、TorchDataを利用する方がよいと考えます。
TorchDataの正式リリースが待たれます。