Pytorch Lightningを使用したBERT文書分類モデルの実装

# はじめに

PytorchでのBERTの分類モデル実装時に、コードの長さと使いまわしのしにくさを感じていたところ、Pytorch Lightningというフレームワークを知ったので実装してみました。

※Pytorch Lightningとは何か？については下記の記事が分かりやすく、参考にさせて頂きました。
[PyTorch 三国志（Ignite・Catalyst・Lightning）]
(https://qiita.com/fam_taro/items/c32e0a21cec5704d9a92)

Google Colaboratoryで実行できるnotebookもgitで公開していますので、よろしければ参考にして頂き、間違っている点などあれば是非ご指摘いただけますと幸いです。
[pytorch_lightning_text_classification.ipynb]
(https://github.com/tchih11/qiita/blob/main/notebooks/pytorch_lightning_text_classification/pytorch_lightning_text_classification.ipynb)

# 環境
Google Colaboratory

# 各種インポート
```python
%%capture
!pip install transformers==3.5.1
!pip install fugashi
!pip install ipadic
!pip install pytorch-lightning==1.1.0
```
```python
import os
import numpy as np
import pandas as pd
import pytorch_lightning as pl
import torch
import torch.optim as optim
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from sklearn.model_selection import train_test_split
from torch import nn
from torch.utils.data import DataLoader, Dataset
from transformers import BertModel
from transformers.tokenization_bert_japanese import BertJapaneseTokenizer
```

# モデリング用データセットの前処理
BERTの分類モデルを作成する際には、前処理として文章をtokenizeしてDataLoaderにする必要があります。今回は自作Datasetと```LightningDataModule```を継承したクラスを使用して前処理を行い、モデリングできる状態にしていきます。
## classの定義
```python
class CreateDataset(Dataset):
    """
    DataFrameを下記のitemを保持するDatasetに変換。
    text(原文)、input_ids(tokenizeされた文章)、attention_mask、labels(ラベル)
    """

    def __init__(self, data, tokenizer, max_token_len):
        self.data = data
        self.tokenizer = tokenizer
        self.max_token_len = max_token_len

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        data_row = self.data.iloc[index]
        text = data_row[TEXT_COLUMN]
        labels = data_row[LABEL_COLUMN]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_token_len,
            padding="max_length",
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt',
        )

        return dict(
            text=text,
            input_ids=encoding["input_ids"].flatten(),
            attention_mask=encoding["attention_mask"].flatten(),
            labels=torch.tensor(labels)
        )

class CreateDataModule(pl.LightningDataModule):
    """
    DataFrameからモデリング時に使用するDataModuleを作成
    """
    def __init__(self, train_df, valid_df, test_df, batch_size=16, max_token_len=512, 
                 pretrained_model='cl-tohoku/bert-base-japanese-char-whole-word-masking'):
        super().__init__()
        self.train_df = train_df
        self.valid_df = valid_df
        self.test_df = test_df
        self.batch_size = batch_size
        self.max_token_len = max_token_len
        self.tokenizer = BertJapaneseTokenizer.from_pretrained(pretrained_model)

    def setup(self):
        self.train_dataset = CreateDataset(self.train_df, self.tokenizer, self.max_token_len)
        self.vaild_dataset = CreateDataset(self.valid_df, self.tokenizer, self.max_token_len)
        self.test_dataset = CreateDataset(self.test_df, self.tokenizer, self.max_token_len)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=os.cpu_count())

    def val_dataloader(self):
        return DataLoader(self.vaild_dataset, batch_size=self.batch_size, num_workers=os.cpu_count())

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, num_workers=os.cpu_count())

```

## 使い方
まずは文章とラベルから成るtrain_df、valid_df、test_df(全てDataFrame)を用意します。今回はcriterionに```nn.CrossEntropyLoss```を使用するので、ラベルは数値型に変換しておきます。

```python
# こんな感じ（データはlivedoorのニュースコーパス）
train_df.head(3)
# > 　　text	category
# > 0	スティーブン・...	1
# > 1	10月末、K-1イ ...	0
>#  2	今日は大みそか...	6
```

あとはDataFrameをCreateDataModuleに渡して```setup()```すれば完成です。

```python
# 用意したDataFrameの文章、ラベルのカラム名
TEXT_COLUMN = "text"
LABEL_COLUMN = "category"

# 作ったDataFrameを渡してsetup
data_module = CreateDataModule(train_df,valid_df,test_df)
data_module.setup()
```

setup後のdata_moduleの中身は下記のような形になっています。

```python
# datasetへのアクセス
item = data_module.train_dataset[0]
print(item["input_ids"])
# > tensor([   2,   27,   26,   26,   70,   25,    6,  120,  356,  186,  337,    7,
# >         1266,  307,  163,  239,  181,    6,   93,   40,  123,   29,   14,   16,
# >          ・・・
# >         2690,   23,   19,  322,  197,   11, 1266,   87,  550,  181,    6,  466,
# >         1823,  577,  693,   11,  142,  125,   47,    3])

print(item["input_ids"].shape)
# > torch.Size([512])
```

```python
# dataloaderのバッチの中身を確認
batch = next(iter(data_module.train_dataloader()))
print(batch["input_ids"])
# > tensor([[   2,   55,   28,  ...,   18,   19,    3],
# >         [   2,  569,   67,  ...,  145,  427,    3],
# >         ...,
# >         [   2,   45,   49,  ...,  163, 1143,    3],
# >         [   2,  672,   11,  ...,    8,  451,    3]])

print(batch["input_ids"].shape)
# > torch.Size([16, 512])
#  → batch_size × max_token_lenになっている
```


# モデリングの実行 & 精度検証
## クラスの定義
データの準備が整ったら、```LightningModule```を継承してモデリング用クラスを作成します。
```def xx_step```で、training/validation/testの各フェーズごとにミニバッチの処理を定義することができる点が特徴です。関数名をフックにして各フェーズで必要となる処理、例えばtraining時の```model.train()```や```loss.backward()```、validation時の```model.valid()```や```torch.no_grad()```といった処理は内部的に行ってくれるため、記述は不要です。全フェーズで必要な```.to(device)```などの処理も内部的にやってくれます。
```__init__```、```forward```を定義するところは普通にpytorchで実装する場合とほぼ変わりません。

```python
class TextClassifier(pl.LightningModule):
    def __init__(self, n_classes: int, n_epochs=None, 
                 pretrained_model='cl-tohoku/bert-base-japanese-char-whole-word-masking'):
        super().__init__()

        # モデルの構造
        self.bert = BertModel.from_pretrained(
            pretrained_model, return_dict=True)
        self.classifier = nn.Linear(self.bert.config.hidden_size, n_classes)
        self.n_epochs = n_epochs
        self.criterion = nn.CrossEntropyLoss()

        # BertLayerモジュールの最後を勾配計算ありに変更
        for param in self.bert.parameters():
            param.requires_grad = False
        for param in self.bert.encoder.layer[-1].parameters():
            param.requires_grad = True

    # 順伝搬
    def forward(self, input_ids, attention_mask, labels=None):
        output = self.bert(input_ids, attention_mask=attention_mask)
        preds = self.classifier(output.pooler_output)
        loss = 0
        if labels is not None:
            loss = self.criterion(preds, labels)
        return loss, preds

    # trainのミニバッチに対して行う処理
    def training_step(self, batch, batch_idx):
        loss, preds = self.forward(input_ids=batch["input_ids"],
                                    attention_mask=batch["attention_mask"],
                                    labels=batch["labels"])
        return {'loss': loss,
                'batch_preds': preds,
                'batch_labels': batch["labels"]}

    # validation、testでもtrain_stepと同じ処理を行う
    def validation_step(self, batch, batch_idx):
        loss, preds = self.forward(input_ids=batch["input_ids"],
                                    attention_mask=batch["attention_mask"],
                                    labels=batch["labels"])
        return {'loss': loss,
                'batch_preds': preds,
                'batch_labels': batch["labels"]}

    def test_step(self, batch, batch_idx):
        loss, preds = self.forward(input_ids=batch["input_ids"],
                                    attention_mask=batch["attention_mask"],
                                    labels=batch["labels"])
        return {'loss': loss,
                'batch_preds': preds,
                'batch_labels': batch["labels"]}

    # epoch終了時にvalidationのlossとaccuracyを記録
    def validation_epoch_end(self, outputs, mode="val"):
        # loss計算
        epoch_preds = torch.cat([x['batch_preds'] for x in outputs])
        epoch_labels = torch.cat([x['batch_labels'] for x in outputs])
        epoch_loss = self.criterion(epoch_preds, epoch_labels)
        self.log(f"{mode}_loss", epoch_loss, logger=True)

        # accuracy計算
        num_correct = (epoch_preds.argmax(dim=1) == epoch_labels).sum().item()
        epoch_accuracy = num_correct / len(epoch_labels)
        self.log(f"{mode}_accuracy", epoch_accuracy, logger=True)

    # testデータのlossとaccuracyを算出（validationの使いまわし）
    def test_epoch_end(self, outputs):
        return self.validation_epoch_end(outputs, "test")

    # optimizerの設定
    def configure_optimizers(self):
        # pretrainされているbert最終層のlrは小さめ、pretrainされていない分類層のlrは大きめに設定
        optimizer = optim.Adam([
            {'params': self.bert.encoder.layer[-1].parameters(), 'lr': 5e-5},
            {'params': self.classifier.parameters(), 'lr': 1e-4}
        ])

        return [optimizer]
```
## 使い方

### インスタンスの作成、及び各種設定
クラスを使用してインスタンスを作成したら、EarlyStoppingとモデルの保存先の設定を行います。

```python
# epoch数
N_EPOCHS = 10

# モデルインスタンスを作成
model = TextClassifier(n_classes=9,n_epochs=N_EPOCHS)

# EarlyStoppingの設定
# 3epochで'val_loss'が0.05以上減少しなければ学習をストップ
early_stop_callback = EarlyStopping(
    monitor='val_loss', 
    min_delta=0.05, 
    patience=3, 
    mode='min')

# モデルの保存先
# epoch数に応じて、「epoch=0.ckpt」のような形で指定したディレクトリに保存される
checkpoint_callback = ModelCheckpoint(
    dirpath="./checkpoints",
    filename='{epoch}',
    verbose=True,
    monitor='val_loss',
    mode='min'
)

# Trainerに設定
trainer = pl.Trainer(max_epochs=N_EPOCHS,
                     gpus=1,
                     progress_bar_refresh_rate=30,
                     callbacks=[checkpoint_callback, early_stop_callback])
```


### 学習の実行
設定が終わったらいよいよ学習です。勝手にループしてくれるので、for文の記述は不要です。

```python
# 学習
trainer.fit(model, data_module)
```
進捗もいい感じに表示してくれます。
<img width="800" alt="pytorch_lightning_modeling_progress" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/989042/fdbcfb45-05d1-a035-e87d-803315dea21b.png">


### テストの実行
一番lossが少なかった時点のモデルを使用してtestデータへの当てはめを行います。

```python
# テスト
result = trainer.test(ckpt_path=checkpoint_callback.best_model_path)
# > -----------------------------------------------------------
# > DATALOADER:0 TEST RESULTS
# > {'test_accuracy': 0.9170632222977566,
# >  'test_loss': tensor(0.2891, device='cuda:0'),
# >  'val_accuracy': 0.9196378041878891, ← この数値は最後のvalidation epochの結果（≠ 最良モデル）
# >  'val_loss': tensor(0.3031, device='cuda:0')}
# > -----------------------------------------------------------

```

### TensorBoard 
モデリング用クラス作成時にself.log()で定義した数値はTensorBoardでも確認可能です。

```python
# Google Colaboratoryの場合
%load_ext tensorboard
%tensorboard --logdir /content/lightning_logs
```
<img width="800" alt="pytorch_lightning_modeling_tensor_board" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/989042/4ab3eb61-0086-637d-d0e4-286ea20b7326.png">

### checkpointからのモデルをロード
下記のような形で、checkpointを指定してモデルをロードすることも可能です。

```python
# 指定したcheckpointのモデルをロード
trained_model = TextClassifier.load_from_checkpoint('./checkpoints/epoch=0.ckpt',n_classes=9)
```

# 感想
普通に実装するよりは格段にスッキリするし、使いまわしもしやすくなったと感じます。また、Early Stoppingなどを手軽に実装できるのも嬉しいですね。まだまだ未開拓な部分もたくさんあるので、徐々に開拓していきたいです。

# 参考

- 公式ドキュメント  
[PYTORCH LIGHTNING DOCUMENTATION]
(https://pytorch-lightning.readthedocs.io/en/latest/)
- Pytorch Lightningとは何か的なところ  
[PyTorch 三国志（Ignite・Catalyst・Lightning）]
(https://qiita.com/fam_taro/items/c32e0a21cec5704d9a92)
- 実装（記事）  
[【実装解説】日本語版BERTでlivedoorニュース分類：Google Colaboratoryで（PyTorch）]
(https://qiita.com/sugulu_Ogawa_ISID/items/697bd03499c1de9cf082)  
[速習 pytorch-lightning: 今すぐ機械学習の実験をしたいそこのキミへ]
(https://qiita.com/yamad07/items/d08098544ea5db9c9674)  
[日本語BERTモデルをPyTorch用に変換してfine-tuningする with torchtext & pytorch-lightning]
(https://radiology-nlp.hatenablog.com/entry/2020/01/18/013039)  
- 実装（youtube動画、全部英語）  
[PyTorch Lightning MasterClass]
(https://www.youtube.com/playlist?list=PLaMu-SDt_RB5NUm67hU2pdE75j6KaIOv2)←公式、再生リスト  
[Fine-Tuning BERT with HuggingFace and PyTorch Lightning for Multilabel Text Classification | Dataset]
(https://www.youtube.com/watch?v=wG2J_MJEjSQ)  
[Fine-Tuning BERT with HuggingFace and PyTorch Lightning for Multilabel Text Classification | Train]
(https://www.youtube.com/watch?v=UJGxCsZgalA&t=2777s)