Я думал естественным способом сделать это было использование обратного вызова с методом 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
