pytorch lightningを使い始めました。
学習ループとか、もろもろラッピングしてもらえてとても便利なのですが、ログ(val_lossやval_accuracy)はどこに残っているのか?という部分が謎でした。
最近(1月)、pytorch lightningが0.6にバージョンアップしデフォルトでTensorBoardLoggerが使われているとのことなので、どんな風に使うかgithubのコードから見てみたいと思います。
まずはMNISTのサンプルスクリプトを動かす
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
import os import torch from torch.nn import functional as F from torch.utils.data import DataLoader from torchvision.datasets import MNIST from torchvision import transforms import pytorch_lightning as pl from pytorch_lightning import Trainer print(pl.__version__) ``` '0.6.0' ``` |
サンプルスクリプトはこちらから持ってきます。
https://github.com/PyTorchLightning/pytorch-lightning
CoolSystemクラスを拝借します。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
class CoolSystem(pl.LightningModule): def __init__(self): super(CoolSystem, self).__init__() # not the best model... self.l1 = torch.nn.Linear(28 * 28, 10) def forward(self, x): return torch.relu(self.l1(x.view(x.size(0), -1))) def training_step(self, batch, batch_idx): # REQUIRED x, y = batch y_hat = self.forward(x) loss = F.cross_entropy(y_hat, y) tensorboard_logs = {'train_loss': loss} return {'loss': loss, 'log': tensorboard_logs} def validation_step(self, batch, batch_idx): # OPTIONAL x, y = batch y_hat = self.forward(x) return {'val_loss': F.cross_entropy(y_hat, y)} def validation_end(self, outputs): # OPTIONAL avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean() tensorboard_logs = {'val_loss': avg_loss} return {'avg_val_loss': avg_loss, 'log': tensorboard_logs} def test_step(self, batch, batch_idx): # OPTIONAL x, y = batch y_hat = self.forward(x) return {'test_loss': F.cross_entropy(y_hat, y)} def test_end(self, outputs): # OPTIONAL avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean() tensorboard_logs = {'test_loss': avg_loss} return {'avg_test_loss': avg_loss, 'log': tensorboard_logs} def configure_optimizers(self): # REQUIRED # can return multiple optimizers and learning_rate schedulers # (LBFGS it is automatically supported, no need for closure function) return torch.optim.Adam(self.parameters(), lr=0.02) @pl.data_loader def train_dataloader(self): # REQUIRED return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32) @pl.data_loader def val_dataloader(self): # OPTIONAL return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32) @pl.data_loader def test_dataloader(self): # OPTIONAL return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32) |
training_step
, validation_end
, test_end
にtensorboad_logsとしてログを残す記述がされていますね。
1 2 3 4 |
model = CoolSystem() trainer = Trainer(max_epochs=10) trainer.fit(model) trainer.test() |
学習を終えたら、ログを確認しに行きましょう!
lightning_logsフォルダを確認
一回学習を走らせるごとに、(デフォルトでは)同じ階層に lightning_ligs
というフォルダができているはずです。
中はこのようになっています。

checkpointsフォルダには、今回の学習ループで、一番lossが小さくなったときのモデルのチェックポイント、ckptファイルが残っています。これを使ってチェックポイント時のモデルを再現することができます。
今回探していたログは、上から2番目の evemts.out.tfevents.~~
ファイルです。
こちらを開きましょう。
terminalで以下のように打ち込みます。
1 |
tensorboard --logdir <ログファイルのあるパス> |
そうすると、localhost:6006の案内がでますので、ここから以下のようにログの確認ができます。
1 2 3 4 5 6 7 8 9 10 |
from pytorch_lightning.logging import TensorBoardLogger # default logger used by trainer logger = TensorBoardLogger( save_dir=os.getcwd(), version=self.slurm_job_id, name='lightning_logs' ) Trainer(logger=logger) |
ログを残す場所やファイル名等の変更はできるようです。
(https://pytorch-lightning.readthedocs.io/en/0.6.0/trainer.html こちらより)
ログを直接残したい
ただいちいちTensorboard開くのが面倒なので、ログ格納リストを追加してそこに放り込んでいこうかなと思います。
CoolSystemクラスを少し修正します。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
class CoolSystem(pl.LightningModule): def __init__(self, val_losses=None): super(CoolSystem, self).__init__() # not the best model... self.l1 = torch.nn.Linear(28 * 28, 10) self.val_losses = val_losses def forward(self, x): return torch.relu(self.l1(x.view(x.size(0), -1))) def training_step(self, batch, batch_idx): # REQUIRED x, y = batch y_hat = self.forward(x) loss = F.cross_entropy(y_hat, y) tensorboard_logs = {'train_loss': loss} return {'loss': loss, 'log': tensorboard_logs} def validation_step(self, batch, batch_idx): # OPTIONAL x, y = batch y_hat = self.forward(x) return {'val_loss': F.cross_entropy(y_hat, y)} def validation_end(self, outputs): # OPTIONAL avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean() tensorboard_logs = {'val_loss': avg_loss} self.val_losses.append(avg_loss) return {'avg_val_loss': avg_loss, 'log': tensorboard_logs} def test_step(self, batch, batch_idx): # OPTIONAL x, y = batch y_hat = self.forward(x) return {'test_loss': F.cross_entropy(y_hat, y)} def test_end(self, outputs): # OPTIONAL avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean() tensorboard_logs = {'test_loss': avg_loss} return {'avg_test_loss': avg_loss, 'log': tensorboard_logs} def configure_optimizers(self): # REQUIRED # can return multiple optimizers and learning_rate schedulers # (LBFGS it is automatically supported, no need for closure function) return torch.optim.Adam(self.parameters(), lr=0.02) @pl.data_loader def train_dataloader(self): # REQUIRED return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32) @pl.data_loader def val_dataloader(self): # OPTIONAL return DataLoader(MNdIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32) @pl.data_loader def test_dataloader(self): # OPTIONAL return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32) |
埋もれてしまってわかりにくいですが、修正箇所は2点で
・def __init__
内にval_lossesを追加
・def validation_end
内で self.val_losses.append(avg_loss)
これで、モデルのインスタンス化の際に、ログを格納するリストを与えてあげればうまくいくはずです。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
log = [] model = CoolSystem(val_losses=log) trainer = Trainer(max_epochs=10, early_stop_callback=False) trainer.fit(model) print(log) ``` [tensor(2.2798), tensor(1.3760), tensor(1.3442), tensor(1.4257), tensor(1.3503), tensor(1.3609), tensor(1.3615), tensor(1.3196), tensor(1.3152), tensor(1.3612), tensor(1.3395)] ``` |
学習を終えて中を確認すると、tensor型で値がリストに入ります。
こちらを可視化します。
1 2 3 4 5 6 7 |
import pandas as pd results = pd.DataFrame({"val_loss": torch.tensor(log).numpy()}) import matplotlib.pyplot as plt %matplotlib inline results.plot() |
うまくいきました。
pandasのデータフレーム型に変換するときは、
・リスト→torch.tensor
・torch.tensor→numpy
に変える点だけ注意が必要です。
ということで、ログの可視化でした。ただ何かしらの実験管理ツールを利用するほうが絶対によさそうですので、そちらも調べてみよう。
最後まで読んでいただきありがとうございました。
追記:wandb でやってみました