こちらの記事で、PyTorchと日本語BERTによる固有表現抽出について、学習済みの分類モデルにBertForTokenClassificationを使うことで実装しました。
PyTorch+BERTによる固有表現抽出(NER)の実装

本記事では自分で好きなネットワークを定義してオリジナルのモデルを作成してみます。

具体的にはBERTのアウトプットをさらにLSTMに通して、最後に全結合層で分類します。
こういった試行錯誤を通じて分類精度を上げることが目的です（結論、ほとんど精度は変わらなかったのですが・・・）。

データローダを作成するところまでと、学習・評価の実装（つまり、モデルの定義部分以外）はPyTorch+BERTによる固有表現抽出(NER)の実装とほぼ同じです（ネットワークを自分で定義する関係で、必要なライブラリが少し増えます。）

コードの全量はgithubにも載せているので、とりあえずコピペで動かしたい場合はこちらをご覧になってください。
https://github.com/nukano0522/pytorch/blob/master/bert_ner/torch_ner_bert%2Blstm.ipynb

1. 必要なライブラリをインストール/インポート

colab上で必要なライブラリをインストール

!pip install transformers==4.21.2 fugashi==1.1.2 ipadic==1.0.0

ライブラリをインポートします。

import os
import json
import unicodedata
import itertools
from tqdm import tqdm
import random
import numpy as np
import pandas as pd
from typing import List, Optional, Tuple, Union
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import BertJapaneseTokenizer, BertModel
from transformers.modeling_outputs import TokenClassifierOutput
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from transformers import BertModel

# 日本語学習済みモデル
MODEL_NAME = 'cl-tohoku/bert-base-japanese-whole-word-masking'

2. データの準備、トークナイザの定義、データローダの作成

こちらはPyTorch+BERTによる固有表現抽出(NER)の実装と同じなので、省略します。
ここまで終わっている前提で、この先を進めます。

3. モデルの定義

3.1. コンフィグの設定

モデルを定義する前に、Bertのコンフィグ情報を取得し、ラベルの分類数を9に設定します。（デフォルトが2になっている）

from transformers import BertConfig
config = BertConfig()

config.num_labels = 9

3.2. モデル定義

モデルを定義します。
トークンの分類（BertForTokenClassification）については、GitHubのコードを参考にしています。
https://github.com/huggingface/transformers/blob/v4.22.1/src/transformers/models/bert/modeling_bert.py#L1713

from torch import nn

model = BertModel.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking", add_pooling_layer=False, output_attentions=True, output_hidden_states=True)

class MyBertForTokenClassification(nn.Module):
    """オリジナルのモデルを定義
    """

    def __init__(self):
        super(MyBertForTokenClassification, self).__init__()
        self.num_labels = config.num_labels

        # BERTモデル
        self.bert = model

        # BiLSTM
        self.lstm = nn.LSTM(input_size=config.hidden_size, hidden_size=config.hidden_size, batch_first=True, bidirectional=True)

        # DropOut
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        
        # 全結合層 ※BiLSTMなので隠れ層のパラメータ数が2倍になっていることを考慮
        self.classifier= nn.Linear(config.hidden_size*2, config.num_labels)

        # 重み初期化処理
        nn.init.normal_(self.classifier.weight, std=0.02)
        nn.init.normal_(self.classifier.bias, 0)


    def forward(
        self,
        input_ids,
        token_type_ids,
        attention_mask,
        labels: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
        ):
      
        return_dict = return_dict if return_dict is not None else config.use_return_dict

        outputs = self.bert(
            input_ids,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask,
            return_dict=return_dict,
        )
        # print(outputs.keys())

        # 最終の隠れ層を取得し、BiLSTMにかける
        lstmout, _ = self.lstm(outputs["last_hidden_state"], None)

        # 全結合層による分類結果
        logits = self.classifier(self.dropout(lstmout))

        loss = None
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        # ロスと分類結果を返す
        return TokenClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

3.3. モデルのロード

# GPU使えるならGPU使う
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 学習済みモデルのロード
model = MyBertForTokenClassification()

model.to(device)

4. モデル学習

学習の処理を定義します。

# 最適化器
optimizer = torch.optim.Adam(params=model.parameters(), lr=2e-5)

# モデルを学習させる関数を作成
def train_model(net, dataloaders_dict, optimizer, num_epochs):

    # GPUが使えるかを確認
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("使用デバイス：", device)
    print('-----start-------')

    # ネットワークをGPUへ
    net.to(device)

    # ネットワークがある程度固定であれば、高速化させる
    torch.backends.cudnn.benchmark = True

    # ミニバッチのサイズ
    batch_size = dataloaders_dict["train"].batch_size

    # epochのループ
    for epoch in range(num_epochs):
        # epochごとの訓練と検証のループ
        for phase in ['train', 'val']:
            if phase == 'train':
                net.train()  # モデルを訓練モードに
            else:
                net.eval()   # モデルを検証モードに

            epoch_loss = 0.0  # epochの損失和
            iteration = 1

            # データローダーからミニバッチを取り出すループ
            for batch in (dataloaders_dict[phase]):
                # batchはTextとLableの辞書型変数

                # GPUが使えるならGPUにデータを送る
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                labels = batch["labels"].to(device)

                # optimizerを初期化
                optimizer.zero_grad()

                # 順伝搬（forward）計算
                with torch.set_grad_enabled(phase == 'train'):

                    # BERTに入力
                    output = net(input_ids=input_ids, 
                                          token_type_ids=None, 
                                          attention_mask=attention_mask, 
                                          labels=labels,
                                          return_dict=True)
                    
                    loss = output[0]

                    # 訓練時はバックプロパゲーション
                    if phase == 'train':
                        loss.backward()
                        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                        optimizer.step()

                        if (iteration % 10 == 0):  # 10iterに1度、lossを表示
                            print(f"イテレーション {iteration} || Loss: {loss:.4f}")

                    iteration += 1

                    # 損失の合計を更新
                    epoch_loss += loss.item() * batch_size

            # epochごとのloss
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)

            print(f"Epoch {epoch+1}/{num_epochs} | phase {phase} |  Loss: {epoch_loss:.4f}")

    return net

実際に学習します。

# 学習・検証を実行
num_epochs = 5
net_trained = train_model(model, dataloaders_dict, optimizer, num_epochs=num_epochs)

5. モデル学習、テストデータを使った推論、モデルの評価

以降はPyTorch+BERTによる固有表現抽出(NER)の実装と同じです。
モデル定義だけを変えたので、基本的に他の処理に影響を与えません。

5.1. 評価結果

評価結果は下記のようになりました。
単純にBertForTokenClassificationを使った場合とほとんど変わらない結果になりました。（むしろ少し下がっている）

5.2. おまけ

さらにBertの最終層だけでなく、最終4層を利用したネットワークでも試してみましたので、
気になる方は下記コードをご覧になってください。
https://github.com/nukano0522/pytorch/blob/master/bert_ner/torch_ner_bert4l%2Blstm.ipynb

5. おわりに

本記事では、TransformersのBertForTokenClassificationをそのまま利用せず、LSTMといった他のネットワークを追加することで精度改善が図れないか試してみました。
結果的に精度改善には至らなかったものの、実際にBertForTokenClassificationのコードを確認しながら進めることで、PyTorchによるディープラーニングの実装について少し理解が深まったのでやってみてよかったです。

参考