Colab + PyTorch Lightning + Comet ML + Bert

Colab + PyTorch Lightning + Comet ML - higepon blog の続き。

目標

Tensorflow 2.0 コンペで自分が書いた training のコードを Colab + PyTorch Lightning + CometML に移植する。移植したことでメリットがあるかを検証する。

Google Drive の利用とその失敗

Colab を利用すると 12時間毎にマシンがリセットされる。つまり巨大な kaggle の dataset や生成物が全て消えてしまう。生成物の例としては train/valid split したものとか modelを save したものとか。これらを毎回生成するのは無駄なので Google Drive 上に置いておく。これをダウンロードすれば良い。というコードを書いたのだがうまく行かなかった。生成にかかる時間はほほゼロになったが Google Drive から生成物をダウンロードする時間が支配的になってしまった。ボツ。

移植

transformers Bert のモデルを TF2.0 で使ったモデルにする。

import torch
from transformers import (
    BertModel,
    BertTokenizer
)

bert_model_name = 'bert-large-uncased-whole-word-masking-finetuned-squad'
tokenizer = BertTokenizer.from_pretrained(bert_model_name)
bert = BertModel.from_pretrained(bert_model_name)

次に LightningModule の Optional な callback をコメントアウトする。validation_step, validation_end, test_step, test_end は Optional である。 forward と training_step を埋めてこんな感じになった。batch が dataloader の返すものかな。

    def training_step(self, batch, batch_nb):
        # batch
        input_ids, attention_mask, token_type_ids, no_answers, y_batch = batch
         
        # fwd
        y_hat = self.forward(input_ids, attention_mask, token_type_ids)
        
        # loss
        loss = loss_fn(y_hat, y_batch, no_answers)
        
        # logs
        tensorboard_logs = {'train_loss': loss}
        experiment.log_metric('train_loss', loss.detach().cpu().numpy(), step=self.global_step, epoch=self.current_epoch, include_context=True)
        return {'loss': loss, 'log': tensorboard_logs}

謎の CUDA エラー

データローダーの実装が終わり end to end で動かしてみると以下のエラー。

RuntimeError: CUDA error: device-side assert triggered

付加情報がない。普通の PyTorch NN なら同じモデルを CPU で動かすと詳細なエラーが分かるのだが transformer Bert はどこかで GPU 前提っぽい。と思ったが CPU 上で動かせた。num_classes が 5 であるべきところが別の config file が使われて 2 になっていた。これが cross entropy loss を計算するところでこけてた。