RL での batch size

Reinforcement LearningWelcome to Spinning Up in Deep RL! — Spinning Up documentation で勉強しながら実装している。とある実装で batch size = 5000 となっていて「値が大きすぎる」と思い、何気なく小さな値に変更した。それをすっかり忘れて試行錯誤しているうちに policy gradient (logprob) が 0.0 になってしまい学習が進まない減少に悩まされた。ログを見て、よくよく考えてみたら logprob が 0 ってことは選択された action の確率が 1 ってことだ。つまり policy はどんな状態でも1つのアクションしか取りようがない状態をモデルが学習してしまっていた。さらに観察すると、これはある episode で agent が右に移動し続けただけで大きな報酬を得てしまったのを学習したのだとわかった。batch size が小さかったせいでこれが大きく聞きすぎてしまったようだ。

OpenAI Gym の MountainCar が難しいという話

Vanilla Policy Gradient を勉強・実装している。CartPole はうまく学習できるのに MountainCar の学習が進まない reward -200 であり続けるという減少にハマった。RL のデバッグは本当に難しく、観察してもよくわからなかった。検索してみるとまさに同じ減少で困っている人が解決方法を紹介していてくれた。結論だけ書くと official env の reward が厳しすぎるので独自の reward を定義しようというものであった。納得。あとから見返してみれば明白なのだが気づけなかった。

Solving Curious case of MountainCar reward problem using OpenAI Gym, Keras, TensorFlow in Python – A Software Engineer's Journal

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 を計算するところでこけてた。

CUDA error: device-side assert triggered

CUDA error: device-side assert triggered のデバッグで苦労した。結論からいうと CUDA error は本当にデバッグが難しい。同じモデルを GPU ではなくて CPU 上で動かせば、もっと親切なエラーメッセージが見られる。CPU で再現しない場合は CUDA_LAUNCH_BLOCKING=1 をセットすること。

Colab + PyTorch Lightning + Comet ML

背景

Kaggle の上位ランカーが PyTorch Lightning について言及していたの試してみる。同様に Comet ML も。Kaggle の試行錯誤を Colab (or Colab Pro) に移行できるかもあわせて検討する。

ToDO

以下淡々と ToDOをこなしていきメモを残す。

Lightning 基礎

  • Lightning の transformers example を Colab 単体で動かす。
  • 上記の dataloader を少ないデータに改造 end to end で素早く回せるようにする。

Lightning training resume

Checkpointing によると2つ方法がある。

1つめ

Trainer の resume_from_checkpoint 引数で checkpoint のパスを指定する。

resume = False
if resume:
  trainer = pl.Trainer(gpus=1,
                                    max_epochs=4,
                                    resume_from_checkpoint='lightning_logs/version_0/checkpoints/_ckpt_epoch_1.ckpt')  
else:
  !rm -rf lightning_logs/  
  trainer = pl.Trainer(gpus=1, max_epochs=2)

trainer.fit(bert_finetuner) 

まず最初の training 時には以下のようなログ。

Epoch 1: 100%|██████████| 4/4 [00:01<00:00,  2.22batch/s, avg_val_acc=0.2, batch_idx=2, gpu=0, loss=1.138, v_num=0, val_loss=1.12]
Epoch 2: 100%|██████████| 4/4 [00:29<00:00,  7.26s/batch, avg_val_acc=0.2, batch_idx=2, gpu=0, loss=1.084, v_num=0, val_loss=1.08]

次に resume すると

Epoch 2: 100%|██████████| 4/4 [00:01<00:00,  2.22batch/s, avg_val_acc=0.6, batch_idx=2, gpu=0, loss=0.992, v_num=1, val_loss=0.916]
Epoch 3: 100%|██████████| 4/4 [00:01<00:00,  2.34batch/s, avg_val_acc=0.6, batch_idx=2, gpu=0, loss=0.952, v_num=1, val_loss=0.844]
Epoch 4: 100%|██████████| 4/4 [00:05<00:00,  1.14s/batch, avg_val_acc=0.8, batch_idx=2, gpu=0, loss=0.911, v_num=1, val_loss=0.751]
Epoch 5: 100%|██████████| 4/4 [00:31<00:00,  7.86s/batch, avg_val_acc=1, batch_idx=2, gpu=0, loss=0.865, v_num=1, val_loss=0.659]

training をresumeできる。ただし Epoch 2 を2回やっているように見える。注意しないといけないのは dataloader は state が resume されるわけではないこと。なので shuffle 前提であるとどこかに書いてあった。

2つめ

experiment version 番号を使う方法。experiment version については、pytorch_lightning.loggers.tensorboard moduleの version 引数に説明がある。version は上記のversion_0 とか version_1 が自動で割り振られるをやめて特定の物を指定する。

コードは

from pytorch_lightning.logging.tensorboard import TensorBoardLogger
bert_finetuner = BertMNLIFinetuner()

resume = False # True

logger = TensorBoardLogger(save_dir='experiments', version=10)
if resume:
  trainer = pl.Trainer(gpus=1, max_epochs=5, logger=logger)
else:
  !rm -rf default/version_10/  
  trainer = pl.Trainer(gpus=1, max_epochs=2, logger=logger)
trainer.fit(bert_finetuner) 

Train 開始。

Epoch 1: 100%|██████████| 4/4 [00:02<00:00,  1.49batch/s, avg_val_acc=0, batch_idx=2, gpu=0, loss=1.122, v_num=10, val_loss=1.26]
Epoch 2: 100%|██████████| 4/4 [00:29<00:00,  7.35s/batch, avg_val_acc=0.2, batch_idx=2, gpu=0, loss=1.032, v_num=10, val_loss=1.18]

Train resume。

INFO:root:model and trainer restored from checkpoint: /content/default/version_10/checkpoints/_ckpt_epoch_1.ckpt
Epoch 2: 100%|██████████| 4/4 [00:01<00:00,  1.54batch/s, avg_val_acc=1, batch_idx=2, gpu=0, loss=0.844, v_num=10, val_loss=0.56]
Epoch 3:  75%|███████▌  | 3/4 [00:01<00:00,  1.54batch/s, avg_val_acc=1, batch_idx=2, gpu=0, loss=0.767, v_num=10, val_loss=0.56]
Epoch 3: 100%|██████████| 4/4 [00:02<00:00,  1.54batch/s, avg_val_acc=1, batch_idx=2, gpu=0, loss=0.767, v_num=10, val_loss=0.446]
Epoch 4: 100%|██████████| 4/4 [00:05<00:00,  1.54batch/s, avg_val_acc=1, batch_idx=2, gpu=0, loss=0.696, v_num=10, val_loss=0.356]
Epoch 5: 100%|██████████| 4/4 [00:31<00:00,  7.81s/batch, avg_val_acc=1, batch_idx=2, gpu=0, loss=0.619, v_num=10, val_loss=0.287]

こちらの方法が自動的に新しい version が振られることなく自然で使いやすいと思う。

Google Drive への保存

Colab はいつセッションが切れるか分からないので Google Drive に保存する。以下のようなコードで Google Drive がマウントされていない場合のみにマウント処理。

from google.colab import drive
from pathlib import Path

def mount_drive_if_necessary():
    drive_path = Path('/content/drive')
    if not drive_path.exists():
      drive.mount(str(drive_path))

mount_drive_if_necessary()

train のコードは

from pytorch_lightning.logging.tensorboard import TensorBoardLogger
bert_finetuner = BertMNLIFinetuner()

resume = True

save_root_path = Path('/content/drive/My Drive/kaggle/tf2.0')
logger = TensorBoardLogger(save_dir= save_root_path / 'logs', version=10, name='simple_bert')
if resume:
  trainer = pl.Trainer(gpus=1, max_epochs=5, logger=logger, default_save_path=save_root_path)
else:
  !rm -rf default/version_10/  
  trainer = pl.Trainer(gpus=1, max_epochs=2, logger=logger, default_save_path=save_root_path)

trainer.fit(bert_finetuner) 

実際の check points と log は以下のように保存される。

/content/drive/My Drive/kaggle/tf2.0
/content/drive/My Drive/kaggle/tf2.0/logs
/content/drive/My Drive/kaggle/tf2.0/logs/simple_bert
/content/drive/My Drive/kaggle/tf2.0/logs/simple_bert/version_10
/content/drive/My Drive/kaggle/tf2.0/logs/simple_bert/version_10/events.out.tfevents.1582691888.8883d33e5953.1505.0
/content/drive/My Drive/kaggle/tf2.0/logs/simple_bert/version_10/meta_tags.csv
/content/drive/My Drive/kaggle/tf2.0/logs/simple_bert/version_10/events.out.tfevents.1582692040.8883d33e5953.1505.1
/content/drive/My Drive/kaggle/tf2.0/simple_bert
/content/drive/My Drive/kaggle/tf2.0/simple_bert/version_10
/content/drive/My Drive/kaggle/tf2.0/simple_bert/version_10/checkpoints
/content/drive/My Drive/kaggle/tf2.0/simple_bert/version_10/checkpoints/_ckpt_epoch_1.ckpt
/content/drive/My Drive/kaggle/tf2.0/simple_bert/version_10/checkpoints/_ckpt_epoch_4.ckpt

ハイパーパラメータを checkpoint と同時に保存する

以下のような警告が出ているので調べる。

UserWarning: Did not find hyperparameters at model.hparams. Saving checkpoint without hyperparameters
  "Did not find hyperparameters at model.hparams. Saving checkpoint without"

hparams が checkpoint と一緒に保存されてうれしいのは load_from_checkpoint で load するとき。Trainer 経由で resume するときはその限りではない気がする。 ちなみに hparams はモデルで

class BertMNLIFinetuner(pl.LightningModule):

    def __init__(self, hparams):
        super(BertMNLIFinetuner, self).__init__()
        
        self.bert = bert
        self.W = nn.Linear(bert.config.hidden_size, 3)
        self.num_classes = 3
        self.hparams = hparams
        self.learning_rate = hparams.learning_rate

...snip...
hparams = Namespace(**{'learning_rate': 2e-05})

bert_finetuner = BertMNLIFinetuner(hparams)

こうすることで hparams が別ファイルとして保存される。

! cat tf2.0/logs/simple_bert/version_10/meta_tags.csv
key,value
learning_rate,2e-05

この hparams と comet.ml などの連携も調べなければいけない。

Comet ML

以下のように Trainer に CometLogger を渡せばOK。オフラインモードはインターネット接続がない環境向けなので使わない。

comet_logger = CometLogger(
    api_key='',
    workspace='higepon',
    project_name="simple-bert-test", # Optional
    experiment_name="my_experiment_name_long7",
    rest_api_key = '')

trainer = pl.Trainer(gpus=1, max_epochs=30, logger=comet_logger, default_save_path=save_root_path)

Code タブ

各 Experiment ごとにソースコードを紐付けることができる。そうすれば各バージョンごとに Diff がとれるようになる。Colab の場合以下のようにソースコードを Comet に送る。

code = ''
for cell_in in In:
  code += cell_in + '\n'

comet_logger.experiment.set_code(code)

f:id:higepon:20200227154801p:plain

Metrics を追加する

デフォルトでは train loss しか Comet で見ることができない。自分の場合は val loss も欲しかったので以下のようにした。

    def validation_end(self, outputs):
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        avg_val_acc = torch.stack([x['val_acc'] for x in outputs]).mean()

        tensorboard_logs = {'val_loss': avg_loss, 'avg_val_acc': avg_val_acc}
        self.logger.experiment.log_metric('val_loss', avg_loss.detach().cpu().numpy(), step=self.global_step, epoch=self.current_epoch, include_context=True)
        return {'avg_val_loss': avg_loss, 'progress_bar': tensorboard_logs}

training の resume

CometLogger は training の resume に対応していないことがわかった(pytorch-lightning/comet.py)。というか TensorboardLogger が resume に対応しているのだね。つまり Lightning と Comet を同時に使いたいなら CometLogger 経由ではなく Comet API を直接叩くのが良さそう。

存在する experiment に追記できてる?

Experiment Overview - Comet.ml よると experiment.get_key(self) を利用して ExistingExperiment を作れる。実際に試してみたらうまく行った。

if resume:
   experiment = comet_ml.ExistingExperiment(api_key='',
                                            previous_experiment=experiment_key,
                                      project_name="simple-bert-test",
                                      workspace='higepon',
                                      ) 
else:
  experiment = comet_ml.Experiment(api_key='',
                                      project_name="simple-bert-test",
                                      workspace='higepon',
                                      )

apex amp を有効にする

Trainer に use_amp=True を渡す。Colab で環境を整えるのはちょっと面倒。

try:
  from apex import amp
except ImportError:
  !pip install https://download.pytorch.org/whl/cu100/torch-1.2.0-cp36-cp36m-manylinux1_x86_64.whl
  !pip install https://download.pytorch.org/whl/cu100/torchvision-0.4.0-cp36-cp36m-manylinux1_x86_64.whl    
  !git clone https://github.com/NVIDIA/apex
  !cd apex;pip3 install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .

! pip -q install pytorch-lightning==0.6.0
! pip -q install transformers==2.1.1
! pip -q install comet_ml

続きます。

GCP + Visual Studio Code 開発環境構築

前提

GCP への ssh 接続。GCP の AI notebooks インスタンスに ssh できなかったので調べた - higepon blog

目標

  • Visual Studio Code のリモート接続を利用して GCP 上のコードを快適に編集。
  • git 周りの操作もそのまま行いたい。
  • Visual Studio Code 自体の習熟。
  • Kaggle 用に directory layout も最適化。

手順

GCP の AI notebooks インスタンスに ssh できなかったので調べた

動機

Visual Studio Code でコードをリモート編集したい。そのためには local Mac からインスタンスssh が必要。

事実

ssh -v

ssh -v  [外部IPアドレス]
OpenSSH_7.9p1, LibreSSL 2.7.3
debug1: Reading configuration data /Users/higepon/.ssh/config
debug1: Reading configuration data /etc/ssh/ssh_config
debug1: /etc/ssh/ssh_config line 48: Applying options for *
debug1: Connecting to [外部IPアドレス] [[外部IPアドレス]] port 22.
debug1: connect to address[外部IPアドレス] port 22: Operation timed out
ssh: connect to host [外部IPアドレス] port 22: Operation timed out

ファイアーウォールの設定

まずヘルプドキュメントを読む。

設定してみよう。"VPC ネットワーク" - "ファイアーウォールルール" を開く。

  1. ファイアウォール ルールの作成
  2. 名前: higepon-kaggle-ssh
  3. ネットワーク:インスタンスのネットワークと同じものを選ぶこと!
  4. 優先度:999
  5. 上り
  6. ターゲットタグ:higepon-kaggle-ssh
  7. ソースIPの範囲:0.0.0.0/0
  8. tcp:22
  9. インスタンス詳細設定
  10. ネットワークタグ:higepon-kaggle-ssh
  11. ssh -v [外部IPアドレス] で無事接続できた!

## 自分用メモ
インスタンス側で ssh 認証鍵を higepon と jupyter 両方で登録しておくこと。

Host kaggle-gcp
    HostName 外部IPアドレス
    User jupyter
    IdentityFile /Users/higepon/.ssh/google_compute_engine
    UserKnownHostsFile=/Users/higepon/.ssh/google_compute_known_hosts
    HostKeyAlias=compute.xxxxx
    IdentitiesOnly=yes
    CheckHostIP=no