Вывод PyTorch Lightning после каждой эпохи ⇐ Python

[Anonymous]

Я использую Pytorch Lightning и после каждой эпохи выполняю выводы на небольшом наборе данных, чтобы получить цифру, которую отслеживаю с помощью веса и смещений.
Я думал естественным способом сделать это было использование обратного вызова с методом on_train_epoch_end, который генерирует график. Последний метод требует выполнения некоторого вывода, поэтому я хотел использовать тренер.predict. Тем не менее, при этом я получаю сообщение об ошибке ниже, так что я думаю, это не тот способ сделать это.
Минимальный воспроизводимый пример:

Код: Выделить всё

import lightning as L
from lightning.pytorch.callbacks import Callback

import torch
from torch.utils.data import DataLoader
from torch import nn, optim

class Model(L.LightningModule):
def __init__(self):
super().__init__()
self.f = nn.Linear(10, 1)

def training_step(self, batch, *args):
out = self(batch)
return out.mean() ** 2

def forward(self, x):
return self.f(x)[:, 0]

def train_dataloader(self):
return DataLoader(torch.randn((100, 10)))

def predict_dataloader(self):
return DataLoader(torch.randn((100, 10)))

def predict_step(self, batch):
return self(batch)

def configure_optimizers(self):
optimizer = optim.Adam(self.parameters(), lr=1e-3)
return optimizer

class CallbackExample(Callback):
def on_train_epoch_end(self, trainer: L.Trainer, model: Model) -> None:
loader = model.predict_dataloader()
trainer.predict(model, loader)

... # save figure to wandb

model = Model()
callback = CallbackExample()
trainer = L.Trainer(max_epochs=2, callbacks=callback, accelerator="mps")

trainer.fit(model)

Код: Выделить всё

File ~/Library/Caches/pypoetry/virtualenvs/novae-ezkWKrh6-py3.9/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py:233, in _LoggerConnector.metrics(self)
231 """This function returns either batch or epoch metrics."""
232 on_step = self._first_loop_iter is not None
--> 233 assert self.trainer._results is not None
234 return self.trainer._results.metrics(on_step)

AssertionError:

Какой способ сделать это наиболее естественный и элегантный?


Подробнее здесь: https://stackoverflow.com/questions/782 ... each-epoch