Repository: Lightning-AI/pytorch-lightning (Python)

Issue:
when the validation_step function returns a type defaultdict, TypeError: first argument must be callable or None occurs
## 🐛 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 #9261 
- Left a comment in https://github.com/PyTorchLightning/pytorch-lightning/pull/9261#discussion_r741130894


### To Reproduce

```python
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

```bash
  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

```bash
* 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
```


Discussion:
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`
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
> <https://github.com/PyTorchLightning/pytorch-lightning/issues/10308#issuecomment-957857971>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ACFDBU5OT3T74C23PX5GROTUKABOBANCNFSM5HGW2SGA>
> .
>

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
>> <https://github.com/PyTorchLightning/pytorch-lightning/issues/10308#issuecomment-957857971>,
>> or unsubscribe
>> <https://github.com/notifications/unsubscribe-auth/ACFDBU5OT3T74C23PX5GROTUKABOBANCNFSM5HGW2SGA>
>> .
>>
>

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:
```python
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:

```python
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.




furthermore i'm seeing issues on lightning=2.2.1 using strategy="ddp" for trainer
Hello!

I am using torchgeo with a self-defined datamodule [`GeoDataModule`](https://github.com/microsoft/torchgeo/blob/ea8306c4fcbfc9dd8940c67133e5e7737ad754c2/torchgeo/datamodules/geo.py#L168). 

`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:

```sh
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](https://github.com/microsoft/torchgeo/blob/ea8306c4fcbfc9dd8940c67133e5e7737ad754c2/torchgeo/datasets/utils.py#L410). In my conda environment I do have 

```py
if isinstance(data, defaultdict):
       return elem_type(data.default_factory, OrderedDict(out))
```

introduced by PRQ #10316. Can someone help me out? 



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b9c808464dadc..eabd9329f0c73 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,89 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 
+## [1.6.0] - 2022-MM-DD
+
+### Added
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+### Changed
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+### Deprecated
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+### Removed
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+### Fixed
+
+- Fixed `apply_to_collection(defaultdict)` ([#10316](https://github.com/PyTorchLightning/pytorch-lightning/issues/10316))
+
+
+-
+
+
+-
+
+
+-
+
+
+-
+
+
 ## [1.5.0] - 2021-11-02
 
 ### Added
@@ -132,7 +215,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for empty `gpus` list to run on CPU ([#10246](https://github.com/PyTorchLightning/pytorch-lightning/pull/10246))
 - Added a warning if multiple batch sizes are found from ambiguous batch ([#10247](https://github.com/PyTorchLightning/pytorch-lightning/pull/10247))
 
-
 ### Changed
 
 - Trainer now raises a `MisconfigurationException` when its methods are called with `ckpt_path="best"` but a checkpoint callback isn't configured ([#9841](https://github.com/PyTorchLightning/pytorch-lightning/pull/9841))
@@ -184,7 +266,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Enabled `on_load_checkpoint` for `LightningDataModule` for all `trainer_fn` ([#10238](https://github.com/PyTorchLightning/pytorch-lightning/pull/10238))
 - Allowed separate config files for parameters with class type when LightningCLI is in `subclass_mode=False` ([#10286](https://github.com/PyTorchLightning/pytorch-lightning/pull/10286))
 
-
 ### Deprecated
 
 - Deprecated Trainer argument `terminate_on_nan` in favor of `detect_anomaly`([#9175](https://github.com/PyTorchLightning/pytorch-lightning/pull/9175))
@@ -220,7 +301,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Deprecated `lr_sch_names` from `LearningRateMonitor` ([#10066](https://github.com/PyTorchLightning/pytorch-lightning/pull/10066))
 - Deprecated `ProgressBar` callback in favor of `TQDMProgressBar` ([#10134](https://github.com/PyTorchLightning/pytorch-lightning/pull/10134))
 
-
 ### Removed
 
 - Removed deprecated `metrics` ([#8586](https://github.com/PyTorchLightning/pytorch-lightning/pull/8586/))
@@ -264,7 +344,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Removed automatic patching of `{train,val,test,predict}_dataloader()` on the `LightningModule` ([#9764](https://github.com/PyTorchLightning/pytorch-lightning/pull/9764))
 - Removed `pytorch_lightning.trainer.connectors.OptimizerConnector` ([#10120](https://github.com/PyTorchLightning/pytorch-lightning/pull/10120))
 
-
 ### Fixed
 
 - Fixed ImageNet evaluation in example ([#10179](https://github.com/PyTorchLightning/pytorch-lightning/pull/10179))
@@ -473,7 +552,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added private `prevent_trainer_and_dataloaders_deepcopy` context manager on the `LightningModule` ([#8472](https://github.com/PyTorchLightning/pytorch-lightning/pull/8472))
 - Added support for providing callables to the Lightning CLI instead of types ([#8400](https://github.com/PyTorchLightning/pytorch-lightning/pull/8400))
 
-
 ### Changed
 
 - Decoupled device parsing logic from Accelerator connector to Trainer ([#8180](https://github.com/PyTorchLightning/pytorch-lightning/pull/8180))
diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py
index 3bd920c2e304b..1e981a0f543e7 100644
--- a/pytorch_lightning/utilities/apply_func.py
+++ b/pytorch_lightning/utilities/apply_func.py
@@ -14,7 +14,7 @@
 import dataclasses
 import operator
 from abc import ABC
-from collections import OrderedDict
+from collections import defaultdict, OrderedDict
 from collections.abc import Mapping, Sequence
 from copy import copy
 from functools import partial
@@ -102,6 +102,8 @@ def apply_to_collection(
             )
             if include_none or v is not None:
                 out.append((k, v))
+        if isinstance(data, defaultdict):
+            return elem_type(data.default_factory, OrderedDict(out))
         return elem_type(OrderedDict(out))
 
     is_namedtuple = _is_namedtuple(data)


Test patch:
diff --git a/tests/utilities/test_apply_func.py b/tests/utilities/test_apply_func.py
index 2c131f96ecc6f..da309f7d22b50 100644
--- a/tests/utilities/test_apply_func.py
+++ b/tests/utilities/test_apply_func.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 import dataclasses
 import numbers
-from collections import namedtuple, OrderedDict
+from collections import defaultdict, namedtuple, OrderedDict
 from typing import List
 
 import numpy as np
@@ -153,6 +153,11 @@ def __init__(self, initial_dict):
     reduced = apply_to_collection(to_reduce, int, lambda x: str(x))
     assert reduced == _CustomCollection({"a": "1", "b": "2", "c": "3"})
 
+    # defaultdict
+    to_reduce = defaultdict(int, {"a": 1, "b": 2, "c": 3})
+    reduced = apply_to_collection(to_reduce, int, lambda x: str(x))
+    assert reduced == defaultdict(int, {"a": "1", "b": "2", "c": "3"})
+
 
 def test_apply_to_collection_include_none():
     to_reduce = [1, 2, 3.4, 5.6, 7, (8, 9.1, {10: 10})]