Repository: Lightning-AI/pytorch-lightning
Pull request: Fix frozen dataclass instance error in `apply_to_collection`
## What does this PR do?


changed the set_attribute function in in the apply_to_collection for recursively applying a function to all elements of a certain dtype in utilities/apply_func.py to Work with frozen dataclass Also now. 



Fixes #10873 

### Does your PR introduce any breaking changes? If yes, please list them.

<None>

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [ ] Did you verify new and **existing tests pass** locally with your changes?
- [ ] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun? Yes

Make sure you had fun coding 🙃


Problem statement:
Frozen Dataclasses as samples produce an error in 1.5+
## 🐛 Bug

When using a frozen dataclass to define a sample, an error is thrown.
```
File ".../python3.9/site-packages/pytorch_lightning/utilities/apply_func.py", line 146, in apply_to_collection
  setattr(result, field_name, v)
File "<string>", line 4, in __setattr__
dataclasses.FrozenInstanceError: cannot assign to field 'x'
```
This only started happening in 1.5+, the snippet below runs fine in 1.4.9.

### To Reproduce

```
import os
from dataclasses import dataclass, fields, is_dataclass
from typing import Any, Optional, Sequence

import pytorch_lightning as pl
import torch
from pytorch_lightning import Trainer
from torch import Tensor
from torch.utils.data import DataLoader, Dataset
from torch.utils.data._utils.collate import default_collate_err_msg_format


@dataclass(frozen=True)
class Sample:
    x: torch.Tensor
    y: torch.Tensor


class my_collate:
    def _collate(self, batch: Sequence[Any]) -> Any:
        elem = batch[0]
        elem_type = type(elem)
        if isinstance(elem, Tensor):
            out = None
            if torch.utils.data.get_worker_info() is not None:
                numel = sum(x.numel() for x in batch)
                storage = elem.storage()._new_shared(numel)
                out = elem.new(storage)
            ndims = elem.dim()
            if (ndims > 0) and ((ndims % 2) == 0):
                return torch.cat(batch, dim=0, out=out)
            return torch.stack(batch, dim=0, out=out)
        elif isinstance(elem, int):
            return torch.tensor(batch)
        elif is_dataclass(elem):
            return elem_type(
                **{
                    field.name: self._collate([getattr(d, field.name) for d in batch])
                    for field in fields(elem)
                }
            )
        raise TypeError(default_collate_err_msg_format.format(elem_type))

    def __call__(self, batch: Sequence[Any]) -> Any:
        return self._collate(batch=batch)


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)
        self.labels = torch.randint(0, 2, (length,))

    def __getitem__(self, index):
        return Sample(x=self.data[index], y=self.labels[index])

    def __len__(self):
        return self.len


class MyDataModule(pl.LightningDataModule):
    def __init__(self, batch_size: int = 2):
        super().__init__()
        self.batch_size = batch_size

    def setup(self, stage: Optional[str] = None):
        self._train = RandomDataset(32, 64)
        self._val = RandomDataset(32, 64)
        self._test = RandomDataset(32, 64)
        self._predict = RandomDataset(32, 64)

    def train_dataloader(self):
        return DataLoader(
            self._train, batch_size=self.batch_size, collate_fn=my_collate()
        )

    def val_dataloader(self):
        return DataLoader(
            self._val, batch_size=self.batch_size, collate_fn=my_collate()
        )

    def test_dataloader(self):
        return DataLoader(
            self._test, batch_size=self.batch_size, collate_fn=my_collate()
        )

    def predict_dataloader(self):
        return DataLoader(
            self._predict, batch_size=self.batch_size, collate_fn=my_collate()
        )


class BoringModel(pl.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: Sample, batch_idx):
        loss = self(batch.x).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch: Sample, batch_idx):
        loss = self(batch.x).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch: Sample, batch_idx):
        loss = self(batch.x).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


if __name__ == "__main__":
    dm = MyDataModule()
    dm.prepare_data()
    dm.setup()
    mnist_train = dm.train_dataloader()
    mnist_val = dm.val_dataloader()

    model = BoringModel()

    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
    )
    trainer.fit(model, mnist_train, mnist_val)

```

### Expected behavior

It would be really nice to be able to continue using frozen dataclasses to define samples. Indexing with `batch.x` is nicer (to me) than `batch[0]`. I also think that a sample _should_ be immutable to save from potential bugs.

### Environment

```
* CUDA:
        - GPU:
        - available:         False
        - version:           None
* Packages:
        - numpy:             1.21.4
        - pyTorch_debug:     False
        - pyTorch_version:   1.10.0
        - pytorch-lightning: 1.5.4
        - tqdm:              4.62.3
* System:
        - OS:                Darwin
        - architecture:
                - 64bit
                - 
        - processor:         i386
        - python:            3.8.12
        - version:           Darwin Kernel Version 21.1.0: Wed Oct 13 17:33:23 PDT 2021; root:xnu-8019.41.5~1/RELEASE_X86_64
```
### Additional context

This might seem like a weird pattern, but if you've got a dataloader that can return 2 types of data, e.g. (x,y) pairs, and in some cases (x,z,y) triplets, where Z is some additional data, being able to use `isinstance` is a bit safer than checking some other attribute (e.g. len).


cc @borda

Hints from issue discussion:
Hey @olliethomas,

Would you be interested in making a contribution to fix this bug? I believe this problem is within the apply_to_collection function.

Best,
T.C
Hello! Can I try and contribute to fixing this bug ? @tchaton 
Regards
Hey @abhi-rf,

Yes, go for it :)

Best
T.C
I have a more general question: If you have a frozen dataclass, isn't it kind of the point not being able to mutate it? So why do you want to work around this? 

If it is just because of 
> I also think that a sample should be immutable to save from potential bugs.

I strongly disagree here as well.

Personally I strongly think that we should not include this hack as this goes around the main purpose of this dataclass feature.

> If you have a frozen dataclass, isn't it kind of the point not being able to mutate it?

Yep, exactly. If I input a collection of Sample, where Sample is a frozen dataclass, then as someone not particularly familiar with the internals of PL, I was not expecting these to be mutated. From my perspective as a user (which caused me view this as a bug), the dataclasses were not being mutated prior to v1.5. If this is not the case, and v1.5+ has instead exposed previously unintended behaviour, then this should just be closed, and I'll update my codebase.

> I also think that a sample should be immutable to save from potential bugs.
> I strongly disagree here as well.

This is probably my own preferences sneaking out. I am generally flippant when it comes to memory use and prefer making almost everything immutable, returning a new instance when an operation is performed. I also have the benefit of not working on anything that's to the scale and complexity of PL though 😄 
`apply_to_collection` gets called when Lightning attempts to move the returned data to the device. If the data is not one of a recognized type, we don't modify the batch. Dataclasses are supported but we should exclude the frozen ones imo. Let's check in our implementation if dataclass is frozen and if it is, skip applying the function. Does that make sense? 
Hey @awaelchli,

I am not entirely sure about this. I totally understand the argument above, but at the same time, as a user, I would assume frozen dataclass means the value won't change but not necessarily the device reference of the value.
IMO, I would expect this to work, but I understand the other view too.
 

@tchaton I would assume it not to be touched at all since it is supposed to be immutable. 

@olliethomas if you want this behavior for now, you could probably work around this by manually defining the `transfer_batch_to_device` hook in your `LightningModule`/`DataModule` like so:

```python
import dataclasses
def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any:
    # batch is a mutable dataclass here
    batch_on_device = super().transfer_batch_to_device(batch, device, dataloader_idx)
    # batch_on_device is still a mutable dataclass
   frozen_batch = MyFrozenDataClass(**dataclasses.asdict(batch_on_device)) # this also has tensors on right device then.
   return frozen_batch
```

The only downside would be that you have to define two separate dataclasses here. 
However, this is what I think to be the best way to achieve this since I strongly advocate for not patching around limitations that have been built into python on purpose.
Hey @justusschock,

Yes, I believe this is a valid argument. 

> Dataclasses are supported but we should exclude the frozen ones imo. Let's check in our implementation if dataclass is frozen and if it is, skip applying the function

However, this might trigger a confusing error message. It would fail as the tensors are left on CPU while they should be on GPU.

IMO, we should fail if the user provides frozen dataclass as ignoring them could lead to unwanted behaviour and confusing error messages. Curious to hear your thoughts @carmocca.
I agree with Justus' comments. I'd say failing is the safest option here.

In any case, using tensors with frozen dataclasses is uncommon, as they are generally meant to be mutated for memory efficiency.

Commit message: fix: remove unused imports

Patch to review:
diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py
index aeefc048fe5f4..eff1e6a12b744 100644
--- a/pytorch_lightning/utilities/apply_func.py
+++ b/pytorch_lightning/utilities/apply_func.py
@@ -144,7 +144,10 @@ def apply_to_collection(
                 )
             if not field_init or (not include_none and v is None):  # retain old value
                 v = getattr(data, field_name)
-            setattr(result, field_name, v)
+            try:
+                setattr(result, field_name, v)
+            except Exception:
+                raise dataclasses.FrozenInstanceError("Cannot apply function to Frozen dataclass instance")
         return result
 
     # data is neither of dtype, nor a collection


Code review comments:
diff_hunk: @@ -144,7 +144,10 @@ def apply_to_collection(
                 )
             if not field_init or (not include_none and v is None):  # retain old value
                 v = getattr(data, field_name)
-            setattr(result, field_name, v)
+            try:
+                setattr(result, field_name, v)
+            except Exception:
path: pytorch_lightning/utilities/apply_func.py
text: @user1:
```suggestion
            except dataclasses.FrozenInstanceError:
```

Don't use generic exceptions