when the validation_step function returns a type defaultdict, TypeError: first argument must be callable or None occurs

[jshin49]

🐛 Bug

• both validation_step and validation_epoch_end are defined through overriding
• when the validation_step function returns a type defaultdict, TypeError: first argument must be callable or None occurs.
• returning types in List, Dict, OrderedDict do not give such error.
• Error occurs from pytorch-lightning>=1.4.6
• Error caused by PR Move tracking epoch end outputs logic to the EvaluationEpochLoop #9261
• Left a comment in Move tracking epoch end outputs logic to the EvaluationEpochLoop #9261 (comment)

To Reproduce

import os

from collections import defaultdict

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)
        return defaultdict(float)

    def validation_epoch_end(self, outputs):
        return outputs

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    test_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        max_epochs=1,
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
    trainer.test(model, dataloaders=test_data)


if __name__ == "__main__":
    run()

Expected behavior

  File "scripts/train.py", line 103, in fine_tune
    trainer.fit(model, dataloaders["train"], dataloaders["dev"])
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 552, in fit
    self._run(model)
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 922, in _run
    self._dispatch()
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 990, in _dispatch
    self.accelerator.start_training(self)
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 92, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 161, in start_training
    self._results = trainer.run_stage()
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1000, in run_stage
    return self._run_train()
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1035, in _run_train
    self._run_sanity_check(self.lightning_module)
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1122, in _run_sanity_check
    self._evaluation_loop.run()
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 111, in run
    self.advance(*args, **kwargs)
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 110, in advance
    dl_outputs = self.epoch_loop.run(
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 111, in run
    self.advance(*args, **kwargs)
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 126, in advance
    output = recursive_detach(output, to_cpu=self.trainer.move_metrics_to_cpu)
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/utilities/memory.py", line 44, in recursive_detach
    return apply_to_collection(in_dict, torch.Tensor, detach_and_move, to_cpu=to_cpu)
  File "/root/.local/share/virtualenvs/zero/lib/python3.8/site-packages/pytorch_lightning/utilities/apply_func.py", line 109, in apply_to_collection
    return elem_type(OrderedDict(out))
TypeError: first argument must be callable or None

Environment

* CUDA:
        - GPU:
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
        - available:         True
        - version:           11.1
* Packages:
        - numpy:             1.21.2
        - pyTorch_debug:     False
        - pyTorch_version:   1.9.1+cu111
        - pytorch-lightning: 1.4.9
        - tqdm:              4.62.3
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - ELF
        - processor:         x86_64
        - python:            3.8.6
        - version:           #40~20.04.1-Ubuntu SMP Wed Jan 6 10:15:55 UTC 2021

[carmocca]

Hi! Thanks for the detailed report.

While this is being fixed, can't you convert the defaultdict to a dict before returning it?

We need to add support and tests for passing defaultdict to apply_to_collection

[jshin49]

Yes. That fixes the problem on my end, but I'm leaving the report in case you guys ever want to fix it, and for the sake of anyone else who encounters the same issue.

…

On Wed, 3 Nov 2021 at 00:44, Carlos Mocholí ***@***.***> wrote: Hi! Thanks for the detailed report. While this is not fixed, can't you convert the defaultdict to a dict before returning it? We need to add support and tests for passing defaultdict to apply_to_collection — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#10308 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACFDBU5OT3T74C23PX5GROTUKABOBANCNFSM5HGW2SGA> .

[jshin49]

The error itself was very hard to debug and even imagine whats going on.

…

On Wed, 3 Nov 2021 at 00:46, Jay Shin ***@***.***> wrote: Yes. That fixes the problem on my end, but I'm leaving the report in case you guys ever want to fix it, and for the sake of anyone else who encounters the same issue. On Wed, 3 Nov 2021 at 00:44, Carlos Mocholí ***@***.***> wrote: > Hi! Thanks for the detailed report. > > While this is not fixed, can't you convert the defaultdict to a dict > before returning it? > > We need to add support and tests for passing defaultdict to > apply_to_collection > > — > You are receiving this because you authored the thread. > Reply to this email directly, view it on GitHub > <#10308 (comment)>, > or unsubscribe > <https://github.com/notifications/unsubscribe-auth/ACFDBU5OT3T74C23PX5GROTUKABOBANCNFSM5HGW2SGA> > . >

[densechen]

It seems not fully fixed this bug currently with lightning==2.0.1.post0.
More specifically, we can write a snip of code to collect data manually:

def lit_collate(batch):
    items = defaultdict(list)

    for item in batch:
        for k, v in item.items():
            items[k].append(v)

    return items

It works well in this case, however, if we do a light modification to upper code:

def lit_collate(batch):
    items = defaultdict(list)

    for item in batch:
        for k, v in item.items():
            items[k].append(v)
    if "x" in items:
        items["x"] = torch.stack(items["x"])
    return items

If "x" exists, the error will raise again.

[dangthatsright]

furthermore i'm seeing issues on lightning=2.2.1 using strategy="ddp" for trainer

[MathiasBaumgartinger]

Hello!

I am using torchgeo with a self-defined datamodule GeoDataModule.

torchgeo==0.6.0
pytorch==2.4
lightning==2.4.0
lightning-utilities==0.11.7

Similarly, my program crashes in the validation_step with the stacktrace:

Traceback (most recent call last):
File "check_slrum.py", line 59, in <module>
    trainer.fit(model=task, datamodule=module)
File "site-packages/lightning/pytorch/trainer/trainer.py", line 538, in fit
    call._call_and_handle_interrupt(
File "site-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
File "site-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
    return function(*args, **kwargs)
File "site-packages/lightning/pytorch/trainer/trainer.py", line 574, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
File "site-packages/lightning/pytorch/trainer/trainer.py", line 981, in _run
    results = self._run_stage()
File "site-packages/lightning/pytorch/trainer/trainer.py", line 1023, in _run_stage
    self._run_sanity_check()
File "site-packages/lightning/pytorch/trainer/trainer.py", line 1052, in _run_sanity_check
    val_loop.run()
File "site-packages/lightning/pytorch/loops/utilities.py", line 178, in _decorator
    return loop_run(self, *args, **kwargs)
File "site-packages/lightning/pytorch/loops/evaluation_loop.py", line 135, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
File "site-packages/lightning/pytorch/loops/evaluation_loop.py", line 396, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_args)
File "site-packages/lightning/pytorch/trainer/call.py", line 319, in _call_strategy_hook
    output = fn(*args, **kwargs)
File "site-packages/lightning/pytorch/strategies/strategy.py", line 410, in validation_step
    return self._forward_redirection(self.model, self.lightning_module, "validation_step", *args, **kwargs)
File "site-packages/lightning/pytorch/strategies/strategy.py", line 640, in __call__
    wrapper_output = wrapper_module(*args, **kwargs)
File "site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
File "site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
File "site-packages/torch/nn/parallel/distributed.py", line 1632, in forward
    inputs, kwargs = self._pre_forward(*inputs, **kwargs)
File "site-packages/torch/nn/parallel/distributed.py", line 1537, in _pre_forward
    moved_inputs, moved_kwargs = _to_kwargs(
File "site-packages/torch/distributed/utils.py", line 266, in _to_kwargs
    _recursive_to(inputs, target_device, use_side_stream_for_tensor_copies)
File "site-packages/torch/distributed/utils.py", line 155, in _recursive_to
    res = to_map(inputs)
File "site-packages/torch/distributed/utils.py", line 145, in to_map
    return list(zip(*map(to_map, obj)))
File "site-packages/torch/distributed/utils.py", line 149, in to_map
    return [type(obj)(i) for i in zip(*map(to_map, obj.items()))]
File "site-packages/torch/distributed/utils.py", line 149, in <listcomp>
    return [type(obj)(i) for i in zip(*map(to_map, obj.items()))]
peError: first argument must be callable or None

The inputs are indeed a defaultdict too. (defaultdict(<class 'list'> .... The creation of the default dict happens in the torchgeo library. In my conda environment I do have

if isinstance(data, defaultdict):
       return elem_type(data.default_factory, OrderedDict(out))

introduced by PRQ #10316. Can someone help me out?