🔰PyTorchでニューラルネットワーク基礎 #37 【DataCollatorWithPadding】

概要

個人的な備忘録を兼ねたPyTorchの基本的な解説とまとめです。今回は、DataLoaderの引数であるcollate_fnに、Hugging Faceが提供するライブラリを組み込む方法を扱います。BERTタイプネットワークへの転移学習を用いたテキスト分類を演習として取り上げ、Hugging Faceライブラリの使い方に焦点を当てて整理していきます¹。自作collate関数を作成する必要がないので、見た目もスッキリ、データ量の大規模化にも対応できるかな🌻

方針

1. できるだけ同じコード進行
2. できるだけ簡潔（細かい内容は割愛）

演習用のファイル

• トークナイザー: unigram_tokenizer_2k.json
  • wikipediaの日本の歴史で学習した2kのトークナイザー
• 学習データ：history_text_label_id.jsonl
  • 歴史の時代区分で分類したテキストデータ
• 事前学習済みモデル：unigram_2k.model
  • 第31回で事前学習したモデル。
• コード: sample_37.ipynb

今回扱う内容

1. BERETタイプモデルのテキスト分類問題へのファインチューニング
2. DataLoaderのcollate_fnにDataCollatorWithPaddingを利用
3. tokenizerにPreTrainedTokenizerFastを利用
4. step単位で学習するために、データ生成の関数を作成

テキスト分類で扱う文は、1文あたりのトークン数（系列長）がそれぞれ異なります。モデルに入力するにはバッチ内で系列長を揃える必要があり、これまでは自作してきました。しかし、HuggingFaceのDataCollatorWithPaddingをDataLoaderのcollate_fnに指定することによってより簡単に実現できます²。

また今回は、学習の進行を epoch 単位ではなく、パラメータ更新の回数である step 単位で管理する方法にも取り組みます。この場合、データセットを何周でも供給し続ける仕組みが必要になります。itertools.cycleが思いつくのですが、これだと1周目に取り出した要素をそのまま記憶して繰り返すため、2周目以降はシャッフルが効かず、まったく同じ順序のバッチが流れ続けてしまいます。そこで、DataLoader の反復を毎回作り直してバッチを無限に供給するジェネレータを自前で用意する必要がありそうです

図：イメージ図³

1. 異なる系列長のテキストデータ

最初に利用するデータを確認してみます。すでに、トークナイザーを使い、文頭に<bos>、文末に<eos>のIDを追加したID列が準備されています。ラベルもID化されています。

利用するライブラリーのimport部分です。

import pandas as pd
import numpy as np
import torch
import torch.nn as nn

# カスタマイズする部分
from torch.utils.data import Dataset, DataLoader

# HuggingFaceのtransformersライブラリ
from transformers import DataCollatorWithPadding

# トークナイザー
# from tokenizers import Tokenizer  # 今回は使わない
from transformers import PreTrainedTokenizerFast

from sklearn.model_selection import train_test_split
#デバイスの選択
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("利用デバイス:", device)

説明メモ

• tokenizerの読み込み方法を変更します。
• HuggingFaceのPreTrainedTokenizerFastを利用します。
• 自作した特殊トークンや語彙数の呼び出し方が変わります。後でまとめて変更します。

データの読み込み部分です。小さなサイズのデータなのでpandasを使います。データ量が増えてきたらHuggingFaceのライブラリを使おうかなぁ

データ読込

pre_trained_model = "./model/unigram_2k.model"                # 事前学習されたモデル
data_filename = "./data/text_label_id.jsonl"                  # 分類問題のデータ
tokenizer_filename = "tokenizer/unigram_tokenizer_2k.json"    # 保存したトークナイザー

df = pd.read_json(data_filename, lines=True)
# 訓練データと検証データに分割
train_data, test_data = train_test_split(df, stratify=df["label"], random_state=55)

データフレームdfは、text、label、ids(文頭にBOS、文末にEOS)を列名にしたデータとなります。

	text	label	ids
438	それ以外にも、財政支出を抑えながら有為な人材を登用する足高の制、漢訳洋書禁輸の緩和や甘藷栽培...	2	[1, 543, 299, 325, 207, 8, 1416, 45, 560, 69, ...
928	行商人は連雀商人と呼ばれた。	4	[1, 137, 359, 61, 11, 437, 1682, 359, 61, 12, ...
111	また、最重要資源となっていた石油も、制海権をなくしつつあることで日本への輸送が困難となってい...	0	[1, 183, 433, 231, 126, 831, 149, 397, 22, 768...

表：利用するデータのサンプル

2. DataCollatorWithPaddingを使う

Datasetクラス、tokenizerの変換、DataLoaderと順番に確認していきます。

2.1 Datasetクラスのカスタマイズ ⁴

Datasetクラスのカスタマイズは、３種類の特殊メソッド（初期化・長さ・データ取得）である__init__()・__len__()・__getitem__を定義すればOKでした。

class SimpleDataset(Dataset):
    def __init__(self, data):
        # (1) データフレームの列データをリストへ変換。
        self.input_ids = data["ids"].tolist()
        self.labels = data["label"].tolist()
   
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        # (2) キーをHuggingFaceのタイプに統一しておきます。
        # 添字でアクセス
        return {
            "input_ids": self.input_ids[idx],
            "labels": self.labels[idx]
        }

dataset = SimpleDataset(train_data)
print(dataset[0])   # dataset[0]で、__getitem__(0)が実行
# {'input_ids': [1, 1417, 45, 5, 929, 878, 143, 152, 116, 5, 324, 699, 39, 1288, 1762, 202, 45, 669, 7, 929, 93, 8, 776, 7, 1648, 121, 189, 5, 162, 462, 715, 137, 18, 6, 2], 
# 'labels': 2}

説明メモ

• 引数のdataはpandasのデータフレーム、読み込むデータ合わせたキーを__init__で利用しています。完全に利用するデータに特化したclass構成です
• (1) リストへ変換
• (2) __getitem__の戻り値をDataCollatorWithPaddingの要求する形に合わせます。

2.2 tokenizerの修正

修正と言っても簡単で、自作したtokenizerのjsonファイルに記載されている特殊トークンをPreTrainedTokenizerFastによってSpecialTokensMixinクラスで定義されている特殊トークン名に置き換えていくだけです⁵。

変換

tokenizer = PreTrainedTokenizerFast(
    tokenizer_file= tokenizer_filename,
    unk_token="<unk>",
    pad_token="<pad>",
    bos_token="<bos>",
    eos_token="<eos>",
    mask_token="<mask>",
)

説明メモ

• unk_tokenやpad_tokenなどがSpecialTokensMixinで用意されている特殊トークン名となります。
• 特殊トークン名に対して、対応する自作のトークンを割り当てます。
• すでにids化されたデータなので、<pad>があれば十分かと

2.3 DataCollatorWithPaddingとDataLoader

HuggingFace謹製のcollatorであるDataCollatorWithPaddingを利用して、バッチ毎のpadding処理を行います。これでcollate関数の自作から開放されます😆

DataCollatorWithPaddingを使う

hf_collator = DataCollatorWithPadding(
    tokenizer=tokenizer,
    padding= "longest")   # longest or True: バッチ内で最長

説明メモ

• HuggingFaceの Data Collator によると、tokenizerは、"PreTrainedTokenizer or PreTrainedTokenizerFast (The tokenizer used for encoding the data.)"と説明されています。すでに変換済みなので問題なしですね🌵

• paddingオプションは、

  • True：paddingする
  • False：paddingしない
  • longest：バッチ内最長・Trueと同じ
  • max_length：max_lengthオプションで指定するサイズに揃える

  の４パターン３種類からの選択となります。デフォルト値はTrueのpaddingありとなります。

DataCollatorWithPaddingを利用することで、DataLoader周りが少しスッキリします⁶。

train_loader = DataLoader(
    dataset,
    batch_size=8,
    collate_fn = hf_collator,
    shuffle=True,
    drop_last=True
)

実際に学習されるバッチごとに系列長が異なることを確認してみました。ほとんど系列長64になるのですが、時々、60や55など異なる長さになっていることが確認できます。うまく利用できているようです😌

確認

print([x["input_ids"].shape[1] for x in train_loader][:16])  # 16個確認
# [64, 64, 64, 64, 60, 64, 64, 64, 64, 64, 64, 64, 64, 64, 60, 64]

3. テキスト分類へ

テキスト分類への転移学習なので、モデルの設定や学習の流れは第35回とほぼ同一になります。DataCollatorWithPaddingの利用によりtokenizer周りの変更、step毎の学習に変更という２点が異なります。

1. tokenizerで属性を指定する部分を修正
2. step毎の学習ループへの変更

3.1 モデルの設定など

基本的にBERTタイプのテキスト分類のファインチューニングなので、文頭<bos>の特徴量を分類問題用Linearへ入力するネットワークモデルとなります。

モデルの設定と学習ループの直前までです。クリックで展開

class ModelConfig:
    def __init__(self, tokenizer):
        # モデル構造
        self.vocab_size = len(tokenizer) # 変更箇所(1) tokenizer.get_vocab_size()
        self.seq_len = 64
        self.d_model = 64
        self.nhead = 4
        self.dim_feedforward = 256
        self.num_layers = 6
        self.dropout = 0.1
        self.out_features = 5
        
        # 特殊トークンID  変更箇所(2)
        self.pad_token_id = tokenizer.pad_token_id #tokenizer.token_to_id("<pad>")
        self.mask_token_id = tokenizer.mask_token_id #tokenizer.token_to_id("<mask>")
        self.bos_token_id = tokenizer.bos_token_id # tokenizer.token_to_id("<bos>")
        self.eos_token_id = tokenizer.eos_token_id # tokenizer.token_to_id("<eos>")
        self.unk_token_id = tokenizer.unk_token_id # tokenizer.token_to_id("<unk>")

        # 特殊トークンのセット
        self.special_tokens = {
            self.pad_token_id,
            self.mask_token_id,
            self.bos_token_id,
            self.eos_token_id,
            self.unk_token_id,
        }
        
        # 通常トークンのリスト special_tokenを除くトークンのリスト（MLMランダム置換用）
        self.normal_tokens = [
            i for i in range(self.vocab_size) 
            if i not in self.special_tokens
        ]
        
        # 学習設定 (今回は利用しないけど使うと便利かも)
        self.batch_size = 512
        self.learning_rate = 0.0001
        self.max_iter = 2000
        self.mask_prob = 0.15
        self.max_grad_norm = 1.0


class DNN(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        
        # 埋め込み層
        self.token_embedding = nn.Embedding(
            num_embeddings=config.vocab_size, 
            embedding_dim=config.d_model,
            padding_idx=config.pad_token_id
        )
        self.pos_embedding = nn.Embedding(num_embeddings=config.seq_len, embedding_dim=config.d_model)
        
        self.layer_norm = nn.LayerNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        
        # Transformer Encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.d_model,
            nhead=config.nhead,
            dim_feedforward=config.dim_feedforward,
            dropout=config.dropout,
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers, enable_nested_tensor=False)
        
        # MLM用の出力層を分類問題用の出力層に付け替えます
        self.classifier = nn.Linear(in_features=config.d_model, out_features=config.out_features)

    
    def forward(self, x):
        # <PAD>マスクの作成
        src_key_padding_mask = (x == self.config.pad_token_id)
        
        # 埋め込み
        tok_emb = self.token_embedding(x)
        pos_emb = self.pos_embedding(torch.arange(x.size(1), device=x.device))
        x = tok_emb + pos_emb.unsqueeze(0)
        
        x = self.layer_norm(x)
        x = self.dropout(x)
        
        # Transformer Encoder
        h = self.transformer_encoder(x, src_key_padding_mask=src_key_padding_mask)

        # 文ベクトルへの Pooling  <BOS>トークン（先頭）に情報を集約
        pooled = h[:, 0, :]  # [batch, d_model]
        
        # 分類 
        y = self.classifier(pooled)  # [batch, num_labels=5]    
        return y

# 事前学習モデルの設定を分類問題用へ更新
checkpoint = torch.load(pre_trained_model, map_location=device)
config = ModelConfig(tokenizer)
config.__dict__.update(checkpoint["config"])

model = DNN(config).to(device)

# 事前学習した重みを分類問題用へコピー
# 事前学習重み（checkpoint側）をフィルタ
pretrained = checkpoint["model_state_dict"]
pretrained_filtered = {
    k: v for k, v in pretrained.items()
    if (not k.startswith("mlm_head."))      # and (not k.startswith("classifier."))
}

# 実際にコピーする部分
# strict=Falseにすることで、ファイルにはないパラメータは読み込まれず、初期状態のまま保持される
# Missing Keys: 重み付けされないパラメータ
info = model.load_state_dict(pretrained_filtered, strict=False)

# 空のリストが表示されるならコピー成功！
print(f"削除されているか確認: {info.unexpected_keys}")

params_to_update = []
for name , param in model.named_parameters():
    param.requires_grad = False
for name, param in model.transformer_encoder.layers[-2:].named_parameters():  # 最終層＋１
    param.requires_grad = True
    params_to_update.append(param)
    print("更新されるパラメータ:", name)
for name, param in model.classifier.named_parameters():
    param.requires_grad = True
    params_to_update.append(param)
    print("更新されるパラメータ:", name) 

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(params_to_update, lr=config.learning_rate)

3.2 step数による学習

最後に学習ループです。step毎のデータの現れ方に応じて2種類考えてみました。

(1) シャッフルされるのは最初の1回
DataLoaderでシャッフルされた最初の形、train_loaderで生成された1 epoch目の並び順を繰り返すタイプとなります。

train_loader = DataLoader(dataset,...,shuffle=True)に対して、itertoolsライブラリのcycleを使うことで簡単に実装できます。

cycleを利用の例

from itertools import cycle
data_iter = cycle(train_loader)

# 1 epoch 目のリストが繰り返される
for step in range(max_iter):
    batch = next(data_iter)
    x = batch["input_ids"].to(device)
    y = model(x)
    ...

itertools.cycleは、1周目に使用した要素を内部リストに保存して再生しているだけで、DataLoader.__iter__ を呼び直さないようです。新しい順列が生成されないことになります。

(2) エポックを一巡するごとに再シャッフルされるタイプ
DataLoaderでシャッフルされ、train_loaderで生成されたepoch毎にもシャッフルを繰り返すタイプ。ステップごとにシャッフルするには、専用の簡易データローダー（ジェネレータ）を自作することで二重ループを隠しながら見通しよく記述することができます。

dataloader作成

# step数で管理するためのデータローダー関数
def infinite_loader(dataloader):
    while True:
        for batch in dataloader:
            yield batch    # (★)

# epochではなく、step数で計測
data_iter = infinite_loader(train_loader)

# 一巡するたびに新しい順列になる
for step in range(max_iter):
    batch = next(data_iter)
    x = batch["input_ids"].to(device)
    y = model(x)
    ...

説明メモ

• for batch in dataloader に入るタイミングでデータの並び順がshuffle=Trueに従って決まります。(★) yield batchで１バッチずつ作られる形です。

• データが一巡 (1 epoch) すると、並び順が新しく引き直され、2周目となります。以降も毎回違う順序になります。

• step数は、変数更新の回数なので、色々細かいことを除けば、epoch数との関係は、ほぼ次のようになります。

  step数 (更新回数) = epoch数 (LOOPの回数) × ミニバッチ分割数

  max_iterは自由に設定するものなので、式とは無関係なのですが、一応、今までとの関係ということです。

step基準で学習ループを行うには、DataLoader を無限に回し続ける仕組みが必要になります。itertools.cycle を使う方法が手軽に見えますが、これは1周目のバッチを保存して再生するだけなので、2周目以降もまったく同じ順序・同じ内容となります。「シャッフルして使いたい！」という意図からすると、実際に使うのは(2)番目のinfinite_loaderを使うタイプかな

tqdmを利用して、ビジュアル面を格好良くして、学習ループを書いてみました。備忘録とはいえ、見え方も正義😆

step数による学習

from tqdm import tqdm

max_iter = config.max_iter # max_iter: 2000

# (1) tqdm使う準備
pbar = tqdm(range(max_iter))
model.train()   # trainモードを明示

# (2) バーを表示しながらループ
for step in pbar:
    # (3) 初期化、予測、損失・精度、微分して変数更新の流れ
    batch = next(data_iter)
    x = batch["input_ids"].to(device)
    t = batch["labels"].to(device)
    optimizer.zero_grad()
    y = model(x)
    loss = criterion(y,t)
    acc = accuracy(y, t)    
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.max_grad_norm)  # 勾配クリップ
    optimizer.step()
    # (4) set_postfixは辞書が引数になる
    pbar.set_postfix({"loss": f"{loss.item():.4f}", "acc":f"{acc:.3f}"})
    # (5) 時々状況をプリント
    if (step+1)%500 == 0:
        tqdm.write(f"{step+1}-step:\tloss:{loss.item():.3f}\tacc:{acc:.3f}")

説明メモ

• (1) tqdmを使い、28%|██▊ | 551/2000 ようなバーを表示させました。単純に見た目の強化、これまでの形（直下）でOKです

for step in range(max_iter):
    batch = next(data_iter)
    x = batch["input_ids"].to(device)
        ...
    print(f"{step+1}: \tloss:{loss.item():.3f}")

• (2) tqdmを使うときのループで、in range(max_iter)をin pbarにするだけです。
• (3) 変数更新の一連の手順、第１回から変わらずですね
• (4) 途中経過の表示部分。tqdm.set_postfix()を利用して表示内容を指定します。引数は辞書形式です。バーの右側に、損失と精度がリアルタイム表示されます。
• (5) 進捗状況を表示をさせてみました。print()ではなくて、tqdm.write()を使うんですね〜

3.3 検証

test_dataに対してもSimpleDatasetとDataLoaderを利用します。面倒なので全部を1つのバッチとして扱い精度計算しました😆

検証

# (1) テストデータ全件を 1バッチ で処理する DataLoader を作成
test_dataset = SimpleDataset(test_data)
test_all_loader = DataLoader(
    test_dataset, 
    batch_size=len(test_dataset),  # 全件を一つのバッチにする
    shuffle=False,
    collate_fn=hf_collator         # DataCollatorWithPaddingのこと
)

# (2) 1回だけループを回してデータを取り出す
test_batch = next(iter(test_all_loader))

# (3) 検証データで推論
model.eval()
with torch.inference_mode():
    x_test = test_batch["input_ids"].to(device)
    t_test = test_batch["labels"].to(device)

    y_test = model(x_test)
    acc = accuracy(y_test, t_test)

print(f"検証精度: {acc}")
# 検証精度: 0.816

step数を調整することで、検証精度は81%前後で落ち着き、結果は第35回と同程度になるはず

個別の精度
列が予測したラベルです。例えば、正解ラベルが昭和（1行目）、江戸時代4個、奈良時代2個、室町時代1個誤判定していると見てください。43/50=0.86 だけ正解しています。江戸時代が苦戦しているようです。

	予測
	昭和	弥生	江戸	奈良	室町
昭和	43	0	4	2	1
弥生	2	41	1	4	2
江戸	7	0	35	5	3
奈良	1	5	4	40	0
室町	1	0	3	1	45

次回

いよいよGPTタイプ (Transformer Decoderタイプ)による文章生成の予定です。

目次

注

1. 第35回と同一の構造になります。当然ながら検証精度も同程度となります。 ↩

2. DataCollatorWithPaddingの引数であるtokenizerは、型がPreTrainedTokenizerBaseに対応していることが要求されます。今回はTokenizerクラスを変換する必要がありますが、事前に対応するPreTrainedTokenizerBaseへ変換しておけば少しだけ楽になります。 ↩

3. AIに描画してもらったのですが、進歩を感じますね〜。イラストはイメージであり、本文とはあまり関係がありません
   画像生成AIが生み出す画像のテイストも、漸次洗練されているのでスライド風イラストを残しておくと後で感慨深くなるのかっ💦 ↩

4. Datasetクラスのカスタマイズの詳細については第33回を参照してください。 ↩

5. tokenizerの学習で利用したtokenizersライブラリもHuggingFaceのライブラリなので変換はとても簡単なはず😉 ↩

6. collate関数については第34回を参照してください。 ↩