Repository: Lightning-AI/pytorch-lightning
Pull request: Fix attribute error for _gpus_arg_default loading checkpoint prior to 1.2.8
## What does this PR do?

Fixes #7031 

I removed a function in #6898 that I should not have removed. 
Old checkpoints (<1.2.8) have this pickled, so in order to load the checkpoint back we need to have access to the source code. 

## Before submitting
- [ ] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [ ] Did you make sure to update the documentation with your changes? (if necessary)
- [ ] 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 update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
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?
Make sure you had fun coding 🙃


Problem statement:
Can't get attribute '_gpus_arg_default'
## 🐛 Bug
when I use the code
```model = model.load_from_checkpoint(ckpt_path)```
It pops out a problem: ```Can't get attribute '_gpus_arg_default' on <module 'pytorch_lightning.utilities.argparse' from '/opt/conda/envs/lasaft/lib/python3.8/site-packages/pytorch_lightning/utilities/argparse.py'>```

### To Reproduce
Very hard to describe the reproduction, because I could not load model.

### Expected behavior
Succeed in loading the model

### Environment

 - PyTorch Version (e.g., 1.0):1.7.1
 - OS (e.g., Linux): Ubuntu 18.04
 - How you installed PyTorch (`conda`, `pip`, source): conda
 - Build command you used (if compiling from source):
 - Python version: 3.8.5
 - CUDA/cuDNN version: 10.2
 - GPU models and configuration: V100
 - Any other relevant information: PL version is 1.3.0rc1

Hints from issue discussion:
I met the same problem when I upgrade from 1.2.2 to 1.2.8 and from 1.2.6 to 1.2.8
And when I debugged it, I found that the “pytorch_lightning.utilities.argparse” in 1.2.8 does not have "_gpus_arg_default", but “pytorch_lightning.utilities.argparse” in 1.2.6 does have "_gpus_arg_default".
I haven't been able to reproduce this issue. 

Taking the boring model and running this:

```python
import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import Trainer, LightningModule


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):
        """
        Testing PL Module

        Use as follows:
        - subclass
        - modify the behavior for what you want

        class TestModel(BaseTestModel):
            def training_step(...):
                # do your own thing

        or:

        model = BaseTestModel()
        model.training_epoch_end = None

        """
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def loss(self, batch, prediction):
        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))

    def step(self, x):
        x = self(x)
        out = torch.nn.functional.mse_loss(x, torch.ones_like(x))
        return out

    def training_step(self, batch, batch_idx):
        output = self(batch)
        loss = self.loss(batch, output)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        output = self(batch)
        loss = self.loss(batch, output)
        return {"x": loss}

    def test_step(self, batch, batch_idx):
        output = self(batch)
        loss = self.loss(batch, output)
        return {"y": loss}

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]

    def train_dataloader(self):
        return DataLoader(RandomDataset(32, 64))

    def val_dataloader(self):
        return DataLoader(RandomDataset(32, 64))

    def test_dataloader(self):
        return DataLoader(RandomDataset(32, 64))


if __name__ == '__main__':
    model = BoringModel()

    trainer = Trainer(
        fast_dev_run=True,
        gpus=-1,
    )
    trainer.fit(model)
    trainer.save_checkpoint('model.pt')

    BoringModel.load_from_checkpoint('model.pt')
```

Works on 1.2.8. I then tried saving a checkpoint using 1.2.6, and then loading it in 1.2.8 as such:

```python 
BoringModel.load_from_checkpoint('model.pt')
```

And this worked fine. Could you give more details as to how this bug appears so I can reproduce @quancs @sun-peach?
I think the problem is that the args were passed into the model and saved as hyper parameters in 1.2.6 but in 1.2.8 we removed the `_gpus_arg_default` function, so pickle can't find the source code for that :( 
Change was made here: #6898 

You can probably fix it by doing right before you load your checkpoint.
**Workaround:**
```python
from pytorch_lightning.utilities import argparse 
setattr(argparse, "_gpus_arg_default", lambda x: 0)
```
my bad. I didn't think of this. shoot.
we should re-introduce this function and keep it there so old checkpoints who have that flaw can still be loaded.

Commit message: no-cover

Patch to review:
diff --git a/pytorch_lightning/utilities/argparse.py b/pytorch_lightning/utilities/argparse.py
index 5f7c280db80e9..3b2ebfbdf01da 100644
--- a/pytorch_lightning/utilities/argparse.py
+++ b/pytorch_lightning/utilities/argparse.py
@@ -286,6 +286,13 @@ def _gpus_allowed_type(x) -> Union[int, str]:
         return int(x)
 
 
+def _gpus_arg_default(x) -> Union[int, str]:  # pragma: no-cover
+    # unused, but here for backward compatibility with old checkpoints that need to be able to
+    # unpickle the function from the checkpoint, as it was not filtered out in versions < 1.2.8
+    # see: https://github.com/PyTorchLightning/pytorch-lightning/pull/6898
+    return _gpus_allowed_type(x)
+
+
 def _int_or_float_type(x) -> Union[int, float]:
     if '.' in str(x):
         return float(x)


Code review comments:
diff_hunk: @@ -286,6 +286,13 @@ def _gpus_allowed_type(x) -> Union[int, str]:
         return int(x)
 
 
+def _gpus_arg_default(x) -> Union[int, str]:  # pragma: no-cover
+    # unused, but here for backward compatibility with old checkpoints that need to be able to
+    # unpickle the function from the checkpoint, as it was not filtered out in versions < 1.2.8
+    # see: https://github.com/PyTorchLightning/pytorch-lightning/pull/6898
+    return _gpus_allowed_type(x)
path: pytorch_lightning/utilities/argparse.py
text: @user1:
Should we just?:

```suggestion
    pass
```

@user2:
means return None?

@user1:
yes

@author:
I'm trying to figure out if that would break something. sigh

In the other PR #6898 , I removed this code
```python
# nothing was passed into the GPUs argument
    if callable(gpus):
        return None
```

So I think when loading the checkpoint, we should directly replace this legacy function with None.


@user1:
@author tests are green and you have the approvals. Do you want to do this or keep it as is?

@author:
Ok, I will test once again that I can load the old checkpoint, but with this change.