Metrics reduction during logging/checkpointing

[justusschock]

🐛 Bug

When logging and checkpointing/early stopping with the metrics like shown in the code below, I get:

Traceback (most recent call last):
  File "pytorch_lightning/trainer/trainer.py", line 521, in train
    self.train_loop.run_training_epoch()
  File "pytorch_lightning/trainer/training_loop.py", line 590, in run_training_epoch
    self.trainer.run_evaluation(test_mode=False)
  File "pytorch_lightning/trainer/trainer.py", line 628, in run_evaluation
    self.evaluation_loop.on_evaluation_end()
  File "pytorch_lightning/trainer/evaluation_loop.py", line 111, in on_evaluation_end
    self.trainer.call_hook('on_validation_end', *args, **kwargs)
  File "pytorch_lightning/trainer/trainer.py", line 887, in call_hook
    trainer_hook(*args, **kwargs)
  File "pytorch_lightning/trainer/callback_hook.py", line 177, in on_validation_end
    callback.on_validation_end(self, self.get_model())
  File "pytorch_lightning/callbacks/early_stopping.py", line 164, in on_validation_end
    self._run_early_stopping_check(trainer, pl_module)
  File "pytorch_lightning/callbacks/early_stopping.py", line 205, in _run_early_stopping_check
    current = torch.tensor(current, device=pl_module.device)
RuntimeError: Could not infer dtype of Accuracy

Please reproduce using the BoringModel and post here

from tests.base.boring_model import BoringModel, RandomDataset
import torch
import pytorch_lightning as pl


class CustomModel(BoringModel):
    def __init__(self, *args, **kwargs):
        super().__init__()

        self.acc = pl.metrics.Accuracy()
    def training_step(self, batch, batch_idx):
        self.acc(torch.rand(1, 3, device=self.device), torch.randint(0, 2, (1,), device=self.device))
        self.log('train_acc', self.acc, on_step=True, on_epoch=True)
        return super().training_step(batch, batch_idx)

    def validation_step(self, batch, batch_idx):
        self.acc(torch.rand(1, 3, device=self.device), torch.randint(0, 2, (1,), device=self.device))
        self.log('val_acc', self.acc, on_step=True, on_epoch=True)
        return super().validation_step(batch, batch_idx)

if __name__ == '__main__':
    early_stop = pl.callbacks.EarlyStopping(monitor='val_acc', mode='max')
    checkpoint = pl.callbacks.ModelCheckpoint(
        monitor='val_acc',
        save_last=True,
        save_top_k=5,
        mode='max',
)

    pl.Trainer(gpus=[0,], max_epochs=20, callbacks=[early_stop, checkpoint]).fit(CustomModel(), torch.utils.data.DataLoader(RandomDataset(32, 500)))

Environment

pl version: 1.1.0

Additional context

Sometimes I also get

  File "pytorch_lightning/trainer/trainer.py", line 470, in fit
    results = self.accelerator_backend.train()
  File "pytorch_lightning/accelerators/gpu_accelerator.py", line 66, in train
    results = self.train_or_test()
  File "pytorch_lightning/accelerators/accelerator.py", line 65, in train_or_test
    results = self.trainer.train()
  File "pytorch_lightning/trainer/trainer.py", line 521, in train
    self.train_loop.run_training_epoch()
  File "pytorch_lightning/trainer/training_loop.py", line 627, in run_training_epoch
    self.run_on_epoch_end_hook(epoch_output)
  File "pytorch_lightning/trainer/training_loop.py", line 856, in run_on_epoch_end_hook
    self.trainer.logger_connector.on_train_epoch_end()
  File "pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py", line 367, in on_train_epoch_end
    self.cached_results.has_batch_loop_finished = True
  File "pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py", line 441, in has_batch_loop_finished
    self.auto_reduce_results_on_epoch_end()
  File "pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py", line 431, in auto_reduce_results_on_epoch_end
    hook_result.auto_reduce_results_on_epoch_end()
  File "pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py", line 217, in auto_reduce_results_on_epoch_end
    tbptt_outs = tbptt_outs[0].__class__.reduce_across_time(tbptt_outs)
  File pytorch_lightning/core/step_result.py", line 566, in reduce_across_time
    value = torch.tensor(value)
RuntimeError: Could not infer dtype of ABCMeta

which I could not reproduce with open code so far, but I assume it to be caused by the same issue.

cc @tchaton

[edenlightning]

lets make sure we add more test around this!

[kendreaditya]

I am having a similar issue. When I try to log a wandb.Table object, the following error is throw. Whats odd is the table is actually displayed on wandb, and when debugging the self.log(f"{key}", wandb.Table(dataframe=metrics[key]) doesn't actually throw an error. However, when I comment out the line, the error is no longer thrown. As per my knowledge the Validation Sanity Check, the first epoch (which includes tracking the wandb.Table), and the first validation epoch all run fine. Right before the 2nd training epoch is where the error lies. The debugger says the error is occurring trainer.fit(model, train_dataloader, validation_dataloader).

  File ".\src\trainer.py", line 56, in <module>
    TrainerSetup()
  File ".\src\trainer.py", line 47, in __init__
    trainer.fit(model, train_dataloader, validation_dataloader)
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 473, in fit
    results = self.accelerator_backend.train()
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\accelerators\gpu_accelerator.py", line 66,
    results = self.train_or_test()
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\accelerators\accelerator.py", line 69, in 
    results = self.trainer.train()
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 524, in train
    self.train_loop.run_training_epoch()
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\trainer\training_loop.py", line 637, in ru
    self.run_on_epoch_end_hook(epoch_output)
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\trainer\training_loop.py", line 868, in ru
    self.trainer.logger_connector.on_train_epoch_end()
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\trainer\connectors\logger_connector\loggernd
    self.cached_results.has_batch_loop_finished = True
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\trainer\connectors\logger_connector\epoch__finished
    self.auto_reduce_results_on_epoch_end()
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\trainer\connectors\logger_connector\epoch_sults_on_epoch_end
    hook_result.auto_reduce_results_on_epoch_end()
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\trainer\connectors\logger_connector\epoch_    tbptt_outs = tbptt_outs[0].__class__.reduce_across_time(tbptt_outs)
  File "C:\ProgramData\Miniconda3\lib\site-packages\pytorch_lightning\core\step_result.py", line 573, in reduce_
    value = torch.tensor(value)
RuntimeError: Could not infer dtype of Table

This logs my metrics and wandb.Table:

def log_step(self, metrics, **kwargs):
    for key in metrics:
        if "stats" not in key:
            self.log(
                f"{key}",
                metrics[key],
                prog_bar=kwargs["prog_bar"],
                on_step=kwargs["on_step"],
                on_epoch=kwargs["on_epoch"])
        else:
            self.log(
                f"{key}",
                wandb.Table(dataframe=metrics[key]),
                prog_bar=False,
                on_step=kwargs["on_step"],
                on_epoch=kwargs["on_epoch"])

Desperate Solution:
I believe the error is caused by torch.tensor(value) in pytorch_lightning\core\step_result.py (as errored above). A solution would be make a class that inherited wandb.Table and a class that can be converted to a torch.tensor (i.e. a pythonic list). Although this is a very bad way of solving the problem, it is the only thing I could get to work without altering the source code.

class CustomTable(wandb.Table):
    def __init__(self, dataframe):
        super().__init__(list(dataframe.values), dataframe=dataframe)
        self._list = list(dataframe.values)

    def __len__(self):
        """List length"""
        return len(self._list)

    def __getitem__(self, ii):
        """Get a list item"""
        return self._list[ii]

    def __str__(self):
        return str(self._list)

Note: this implementation only works with pandas.DataFrame as argument.