pytorch-lightningでログを可視化したい【機械学習】

pytorch lightningを使い始めました。

学習ループとか、もろもろラッピングしてもらえてとても便利なのですが、ログ（val_lossやval_accuracy）はどこに残っているのか？という部分が謎でした。

最近（1月）、pytorch lightningが0.6にバージョンアップしデフォルトでTensorBoardLoggerが使われているとのことなので、どんな風に使うかgithubのコードから見てみたいと思います。

まずはMNISTのサンプルスクリプトを動かす

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'
```

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クラスを拝借します。

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)

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としてログを残す記述がされていますね。

model = CoolSystem()
trainer = Trainer(max_epochs=10)
trainer.fit(model)
trainer.test()

model = CoolSystem()

trainer = Trainer(max_epochs=10)

trainer.fit(model)

trainer.test()

学習を終えたら、ログを確認しに行きましょう！

lightning_logsフォルダを確認

一回学習を走らせるごとに、（デフォルトでは）同じ階層に lightning_ligs というフォルダができているはずです。

中はこのようになっています。

checkpointsフォルダには、今回の学習ループで、一番lossが小さくなったときのモデルのチェックポイント、ckptファイルが残っています。これを使ってチェックポイント時のモデルを再現することができます。

今回探していたログは、上から2番目の evemts.out.tfevents.~~ ファイルです。

こちらを開きましょう。

terminalで以下のように打ち込みます。

tensorboard --logdir <ログファイルのあるパス>

1	tensorboard --logdir <ログファイルのあるパス>

そうすると、localhost:6006の案内がでますので、ここから以下のようにログの確認ができます。

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)

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クラスを少し修正します。

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(MNｄIST(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)

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(MNｄIST(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)

これで、モデルのインスタンス化の際に、ログを格納するリストを与えてあげればうまくいくはずです。

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)]
```

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型で値がリストに入ります。

こちらを可視化します。

import pandas as pd
results = pd.DataFrame({"val_loss": torch.tensor(log).numpy()})

import matplotlib.pyplot as plt
%matplotlib inline
results.plot()