概要

個人的な備忘録を兼ねたPyTorchの基本的な解説とまとめです。Transformer Encoderタイプ（BERTタイプ）のMLM事前学習に向けて学んだことをまとめてみました。今回は、系列長の異なるID列データをそのまま扱い、テキスト分類を行う演習を行ってみたいと思います。事前にID列を等長化する必要がないので、少しだけ汎用性があるように感じます🌸

扱いたいテーマ

Datasetクラスのカスタマイズ方法（第33回）
DataLoaderのcollate_fnを使ってみる（第34回）
可変長データでの文章分類　（今回・第35回）
可変長データでのBERTタイプの事前学習（次回）
HuggingFaceのライブラリーとつなげてみる　（予定）

方針

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

演習用のファイル

トークナイザー: unigram_tokenizer_2k.json
- wikipediaの日本の歴史で学習した2kのトークナイザー
学習データ：history_text_label_id.jsonl
- 第32回で利用したテキスト分類データ、歴史の時代区分で分類したテキストデータ
保存モデル：unigram_2k.model
- 第31回で事前学習したモデルのこと
コード: sample_35.ipynb

テキスト分類を例にカスタム版Datasetクラスとのシンプルなcollate_fnを利用したDataLoaderを使って演習を行ってみました。基本的な構造は第22回のTransformerでの分類と同一です¹。ファインチューニングという観点では第32回と同様になります。

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

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

データの読み込み

import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from tokenizers import Tokenizer
from sklearn.model_selection import train_test_split

# カスタマイズするときに利用するライブラリ
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence



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

# 精度を計算する関数
def accuracy(y, t):
    _, argmax_list = torch.max(y, dim=1)
    accuracy = sum(argmax_list == t).item()/len(t)
    return accuracy


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

# トークナイザーの読み込み
tokenizer = Tokenizer.from_file(tokenizer_filename)
# データの読み込み
df = pd.read_json(data_filename, lines=True)
train_data, test_data = train_test_split(df, stratify=df["label"], random_state=55)

表１：利用するデータ（train_data）のサンプル

	text	label	ids
468	寛政の改革ではこれまでの収奪一辺倒だった政策を改め、民を救うための政治へと断行した。	2	[1, 1417, 45, 5, 929, 878, 143, 152, 116, 5, 3...
226	海洋リザーバ効果の影響や不安定な年輪年代測定などの問題もあり、正確な時期は確定していない。	1	[1, 196, 1013, 1389, 1328, 507, 3, 1540, 680, ...
646	2月11日、鋳銭をつかさどる催鋳銭司がおかれ、2月15日、平城遷都の詔が出された。	3	[1, 68, 250, 43, 43, 262, 8, 1016, 127, 7, 145...

表だとわかりにくいのですが、ids列は文章を<pad>で等長化していない、系列長が異なる（可変長）データとなっています。

等長化されていないのでTensorDataset・DataLoaderクラスのコンボは利用できません。次の２つで対応していきます。

Datasetクラスを修正（第33回）
DataLoaderのcollate_fn関数でミニバッチ毎にpaddingして等長化（第34回）

2. カスタムDatasetクラスとDataLoaderクラスのcollate_fn

2.1 カスタムDatasetクラス

Datasetクラスをカスタマイズするのって難しそうしかし、第33回で少し調べた結果、自分が想像していたよりも簡単でした。

カスタム版Datasetクラスを作成

3種類の特殊メソッド__init__()、__len__()、__getitem__()を準備する

テキスト分類で利用するデータをpandasのデータフレームとしています。利用するのは、ids列とlabel列となります。この点を考慮して、次のようなカスタムDatasetクラスを作成してみました。

SimpleDataset (カスタム版Datasetクラス)

class SimpleDataset(Dataset):
    # (1) 対応する列をリストへ
    def __init__(self, data):
        self.ids = data["ids"].tolist()
        self.label = data["label"].tolist()
    
    # (2)
    def __len__(self):
        return len(self.label)
    
    # (3) リストのindex番号でアクセス
    def __getitem__(self, idx):
        return {"ids":   self.ids[idx],
                "label": self.label[idx]}

説明メモ

(1) __init__()：入力されるデータフレームをリストへ変換、text列は利用しません。
(2) __len__()：データのサイズ。ラベル側で数をカウントしてみました。
(3) __getitem__()：ids、labelともにリストとなっているので、リストのindex番号で対応できます。

SimpleDatasetクラスとtrain_dataを利用して0番目の値を取得してみました。

確認

dataset = SimpleDataset(train_data)
print(dataset[0])   # dataset[0]でdataset.__getitem__(0)のこと
# {'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], 'label': 2}

うまく動作しているようです🌵

DataLoaderのcollate_fnの引数は、dataset[i]のリストでした。今回は次の形の辞書形式のリスト

[{'ids':[1,...,2], 'label': 0}, {'ids': [1,...,2], 'label': 1},... ]

がDataLoaderのcollate_fnの引数となります。

2.2 collate_fnの定義

DataLoaderで利用するcolalte関数（collate_fn）を定義します。これが最後のカスタマイズ部分となります🍀

collate_fnを設計

バッチ内の最長の系列長に合わせて、各idsに<pad>IDの 0 を挿入してバッチ内で等長化
pytorchのtorch.nn.utils.rnn.pad_sequenceを利用して<pad>を挿入に対応します😆

pad_sequenceを利用するタイプ

# attention maskは自作しないといけないのがポイント
def padding_collate_fn(batch):
    # (1) batchはデータフレーム、tensor化する
    ids_list = [torch.tensor(item["ids"], dtype=torch.long) for item in batch]
    labels = [item["label"] for item in batch]
    
    # (2) <pad>の挿入
    padded_ids = pad_sequence(ids_list, batch_first=True, padding_value=0)
    attention_masks = (padded_ids != 0).long()   # 今回使わないけどpad部分をマスクするattention_maskも作れるぞ〜
  
    # (3) DataLoaderの出力
    return {
        "ids": padded_ids,
        "attention_mask": attention_masks,
        "label": torch.tensor(labels, dtype=torch.long)
    }

説明メモ

(1) 引数のbatchは辞書形式のリスト

[{'ids':[1,...,2], 'label': 0}, {'ids': [1,...,2], 'label': 1},... ]

となります。idsキーとlabelキー毎に分けて、リスト化します。
ids列はpad_seqenceを使うので、変数をtorch.tensorにする必要があるので注意
["ids"]部分はtorch.tensorのリストにします。["label"]はそのままリスト化
(2) pad_sequenceを利用して、batchごとに<pad>の自動挿入。<pad>のIDを「0」としています。padded_idsの要素が0以外だと1という仕組みで、<pad>とそれ以外のマスクも作成してみました²。今回は第31回のモデルを援用するためモデル側でpad用のマスクを作成する方法です。
(3) 出力部分、すべてtorch.tensorとしました。

出力を確認
batch_size=2で出力を確認してみましょう。

train_loader = DataLoader(
    dataset,                            # dataset = SimpleDataset(train_data)
    batch_size=2,                       # 
    collate_fn = padding_collate_fn,    # collate_fnの設定
    shuffle = True,                     # シャッフル
    drop_last=True                      # batch_size未満の分割時には削除
)
# 実際に学習されるバッチごとに系列長が異なることを確認
print([x["ids"].shape[1] for x in train_loader])
# [41, 64, 64, 45, 57, 64, 64, 63, 15,...]
# 具体的な出力の確認
print(next(iter(train_loader)))

表示がわかりにくいのですが、<pad>の挿入後のidsとマスク、ラベルの3種類が出力されています。等長化されていることも確認できます。

next(iter(train_loader))の結果

{'ids': tensor([[1, 1417, 45, 5,  929,  878,  143,  152,  116, 5, 324,  699, 9, 1288, 1762,  202,  45,  669,  7,  929,  93,  8, 776, 7,
                1648,  121,  189,    5,  162,  462,  715,  137,   18, 6, 2, 0, 0, 0, 0, 0, 0],
               [1,  196, 1013, 1389, 1328,  507,  3, 1540,  680,  5, 1737, 1743, 16,  234,  395,   31,   26, 1338,  26,  191, 1699,  134,  205,  667,
               1742,   13,  286,    8,   90, 1867,  31,  467,   80,   11, 1867,  134, 34, 24, 130, 6, 2]]),
 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]),
 'label': tensor([2, 1])}

3. テキスト分類してみた

第31回で事前学習したBERTタイプのモデルを転移学習する形でテキスト分類を行ってみます。コードがやや長くなるので折りたたみで行きます😆

ネットワークの構造はID化された文 (ids) を入力データとします。埋め込み層、Transformer Encoder層と経由、文頭の<bos>の特徴量のみを使い、分類アダプターである全結合層へという流れになります。

データ読込部分に、SimpleDatasetとpadding_collate_fnを使い、異なる長さのidsでも学習できるようになっています。

図１：モデル構造（分類ネットワーク）

コードの全体像となります。クリックで展開

import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from tokenizers import Tokenizer

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

from sklearn.model_selection import train_test_split

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

# 精度を計算する関数
def accuracy(y, t):
    _, argmax_list = torch.max(y, dim=1)
    accuracy = sum(argmax_list == t).item()/len(t)
    return accuracy

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

### 【１】 データの読み込み
df = pd.read_json(data_filename, lines=True)
tokenizer = Tokenizer.from_file(tokenizer_filename)

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

# カスタム版Dataset
class SimpleDataset(Dataset):
    def __init__(self, data):
        self.ids = data["ids"].tolist()
        self.label = data["label"].tolist()
   
    def __len__(self):
        return len(self.label)
    
    # 入力するデータのタイプによって適宜修正
    def __getitem__(self, idx):
        return {"ids":   self.ids[idx],
                "label": self.label[idx]}

# collate_fn関数の定義（pad_sequenceを利用するタイプ）
def padding_collate_fn(batch):
    # (1) batchはデータフレーム、tensor化する
    ids_list = [torch.tensor(item["ids"], dtype=torch.long) for item in batch]
    labels = [item["label"] for item in batch]
    
    # (2) <pad>の挿入
    padded_ids = pad_sequence(ids_list, batch_first=True, padding_value=0)
    attention_masks = (padded_ids != 0).long()   # 今回使わないけどpad部分をマスクするattention_maskも作れるぞ〜
  
    # (3) DataLoaderの出力
    return {
        "ids": padded_ids,
        "attention_mask": attention_masks,
        "label": torch.tensor(labels, dtype=torch.long)
    }


# 学習データのセットアップ
dataset = SimpleDataset(train_data)
train_loader = DataLoader(
    dataset,
    batch_size=8,
    collate_fn = padding_collate_fn,    # collate_fnの設定
    shuffle=True,
    drop_last=True
)

### 【２】 モデルの定義と作成

# --- モデルの設定
# 第29回〜第32回までの内容を参照
# モデル変数の設定
class ModelConfig:
    def __init__(self, tokenizer):
        # モデル構造
        self.vocab_size = tokenizer.get_vocab_size()
        self.d_model = 64
        self.seq_len = 64
        self.nhead = 4
        self.dim_feedforward = 256
        self.num_layers = 6
        self.dropout = 0.1
        self.out_features = 5
        
        # 特殊トークンID
        self.pad_token_id = tokenizer.token_to_id("<pad>")
        self.mask_token_id = tokenizer.token_to_id("<mask>")
        self.bos_token_id = tokenizer.token_to_id("<bos>")
        self.eos_token_id = tokenizer.token_to_id("<eos>")
        self.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.num_epochs = 100
        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)
        
        # --- 事前学習したモデルから置き換わる部分
        # 分類アダプター
        self.classifier = nn.Linear(in_features=config.d_model, out_features=config.out_features)

    
    def forward(self, x):
        # マスクの作成
        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)

        # <BOS>トークン（先頭）に情報が集約されるように学習される予定
        pooled = h[:, 0, :]  # [batch, d_model]
        
        # 分類 
        y = self.classifier(pooled)  # [batch, num_labels=5]    
        return y

# --- 学習済みモデルの読み込み
# 第32回と同様の手順　学習済みモデルを読み込んで必要な部分だけコピーする
# 事前学習モデルの設定を分類問題用へ更新
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."))
}

# 実際にコピーする部分
# 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=1e-3)


### 【４】 変数更新ループ
LOOP = 30
model.train()
for epoch in range(LOOP):
    total_loss = 0
    total_acc = 0
    num_batches = 0
    
    for data in train_loader:
        optimizer.zero_grad()

        # データを取り出してデバイスに送る
        x = data["ids"].to(device)    # (batch_size, seq_length) seq_lengthは学習中のバッチごと異なる
        t = data["label"].to(device)  # (batc_size)
    
        y = model(x)
  
        loss = criterion(y,t)
        acc = accuracy(y, t)
        
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        total_loss += loss.item()
        total_acc += acc
        num_batches += 1

    if (epoch+1)%10 == 0:
        print(f"{epoch+1}:\tloss:{total_loss/num_batches:.3f}\tacc:{total_acc/num_batches:.3f}")


### 【５】検証
# 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=padding_collate_fn
)

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

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

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

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

ちょっとした考察

学習時の精度について。30エポックくらいで学習精度はおよそ95%くらいまで上昇
ミニバッチ化したほうが学習回数が少ない！
ミニバッチ毎にoptimizer.step()で変数を更新するから、変数更新の回数が増えるのが要因か？学習の回数（epoch数）よりも、変数更新の回数（step数）がどうも重要っぽい気がする🤔　次回、step数に注目して学習ループを作成してみよう
検証精度は、81%前後で落ち着くと思います（第32回と同程度かそれ以上）

個別の精度
うまく表にできなかったのですが、列が予測したラベルです。例えば、正解ラベルが昭和（1行目）、江戸時代2個、奈良時代3個誤判定していると見てください。42/50=0.84 だけ正解しています。

	予測
	昭和	弥生	江戸	奈良	室町
昭和	42	0	2	3	3
弥生	4	40	0	3	3
江戸	5	0	37	6	2
奈良	2	3	2	42	1
室町	3	0	1	1	45

江戸時代が正解だけど、昭和と予測（５個）だったり、奈良時代はいろいろ誤判定しているように見えます。まあ適当に作成したデータセットなのでこんな感じでOKでしょう。

目次ページ

注

コードの書き方が変わっただけで、内容面ではあまり進んでいない ↩
<pad>のidはtokenizerから取得する形にしたほうが汎用的かと思います😆　こんなこともできるんだ〜というサンプルになります ↩