Repository: Lightning-AI/pytorch-lightning
Pull request: update batch size in LightningModule.datamodule when auto scaling batch size

## What does this PR do?


Fixes #3233 

the new test fails on master

# Before submitting

- [x] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [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? Otherwise, we ask you to create a separate PR for every change.
- [x] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests? 
- [x] Did you verify new and existing tests pass locally with your changes?
- [x] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in 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.     
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
auto_scale_batch_size not working with datamodule
## 🐛 Bug

The `Trainer` expects the `LightningModule` to have `self.batch_size` (see `scale_batch_size()` in training_tricks.py). However, if one is using the new `LightningDataModule`, that should be the class with `self.batch_size` defined.

### To Reproduce

```python
assert hasattr(lightning_data_module, "batch_size")
trainer = Trainer(auto_scale_batch_size=True)
trainer.fit(lightning_module, datamodule=lightning_data_module)
```

```bash
pytorch_lightning.utilities.exceptions.MisconfigurationException: Field batch_size not found in both `model` and `model.hparams`
```

### Expected behavior

`auto_scale_batch_size` should work using `LightningDataModule`

### Environment

```
* Packages:
	- numpy:             1.18.5
	- pyTorch_debug:     False
	- pyTorch_version:   1.6.0
	- pytorch-lightning: 0.9.1rc1
	- tensorboard:       2.2.0
	- tqdm:              4.48.2
```

Hints from issue discussion:
I knew this bug report will eventually come :) same will happen for auto_lr_find. 
We need to generalize our `lightning_hasattr` and `getattr` helper functions to include the datamodule. In total, we have 
```
model.attribute_name
model.hparams.attribute_name
model.dm.attribute_name
```
all of these should be considered by both lr_find and auto_scale_batch_size
@awaelchli  So currently there is no solution to this issue? 
no, you need to set it manually, sorry. 
But I'll try to make a PR this weekend if not today. Should be a relatively easy fix unless I missed something. 
@awaelchli I completely agree that it should be a simple fix, since all the attribute getting/setting is handled by our own `lightning_getattr` ect functions.

Commit message: docs

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d18fac2668279..d3a862c624709 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed `GpuUsageLogger` to work on different platforms ([#3008](https://github.com/PyTorchLightning/pytorch-lightning/pull/3008))
 
+- Fixed setting batch size in `LightningModule.datamodule` when using `auto_scale_batch_size` ([#3266](https://github.com/PyTorchLightning/pytorch-lightning/pull/3266))
+
 ## [0.9.0] - YYYY-MM-DD
 
 ### Added
diff --git a/pytorch_lightning/trainer/training_tricks.py b/pytorch_lightning/trainer/training_tricks.py
index 32d0c59434c7a..7c982241626dd 100644
--- a/pytorch_lightning/trainer/training_tricks.py
+++ b/pytorch_lightning/trainer/training_tricks.py
@@ -259,7 +259,9 @@ def _adjust_batch_size(trainer,
                        desc: str = None):
     """ Function for adjusting the batch size. It is expected that the user
         has provided a model that has a hparam field called `batch_size` i.e.
-        `model.hparams.batch_size` should exist.
+        `model.hparams.batch_size` should exist. Additionally there can be a
+        datamodule attached to either Trainer or model, in which case the attribute
+        also gets updated when present.
 
     Args:
         trainer: instance of pytorch_lightning.Trainer
@@ -277,16 +279,12 @@ def _adjust_batch_size(trainer,
     """
     model = trainer.get_model()
     batch_size = lightning_getattr(model, batch_arg_name)
-    if value:
-        lightning_setattr(model, batch_arg_name, value)
-        new_size = value
-        if desc:
-            log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
-    else:
-        new_size = int(batch_size * factor)
-        if desc:
-            log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
-        lightning_setattr(model, batch_arg_name, new_size)
+    new_size = value if value is not None else int(batch_size * factor)
+    if desc:
+        log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
+    lightning_setattr(model, batch_arg_name, new_size)
+    if trainer.datamodule is not None and hasattr(trainer.datamodule, batch_arg_name):
+        setattr(trainer.datamodule, batch_arg_name, new_size)
     return new_size
 
 
diff --git a/pytorch_lightning/utilities/parsing.py b/pytorch_lightning/utilities/parsing.py
index 6f57e63e48fc9..8b9cc0e0479b8 100644
--- a/pytorch_lightning/utilities/parsing.py
+++ b/pytorch_lightning/utilities/parsing.py
@@ -175,8 +175,8 @@ def __repr__(self):
 
 
 def lightning_hasattr(model, attribute):
-    """ Special hasattr for lightning. Checks for attribute in model namespace
-        and the old hparams namespace/dict """
+    """ Special hasattr for lightning. Checks for attribute in model namespace,
+        the old hparams namespace/dict, and the datamodule. """
     # Check if attribute in model
     if hasattr(model, attribute):
         attr = True
@@ -186,6 +186,8 @@ def lightning_hasattr(model, attribute):
             attr = attribute in model.hparams
         else:
             attr = hasattr(model.hparams, attribute)
+    elif hasattr(model.datamodule, attribute):
+        attr = True
     else:
         attr = False
 
@@ -193,8 +195,8 @@ def lightning_hasattr(model, attribute):
 
 
 def lightning_getattr(model, attribute):
-    """ Special getattr for lightning. Checks for attribute in model namespace
-        and the old hparams namespace/dict """
+    """ Special getattr for lightning. Checks for attribute in model namespace,
+        the old hparams namespace/dict, and the datamodule. """
     # Check if attribute in model
     if hasattr(model, attribute):
         attr = getattr(model, attribute)
@@ -204,24 +206,36 @@ def lightning_getattr(model, attribute):
             attr = model.hparams[attribute]
         else:
             attr = getattr(model.hparams, attribute)
+    elif hasattr(model.datamodule, attribute):
+        attr = getattr(model.datamodule, attribute)
     else:
-        raise ValueError(f'{attribute} is not stored in the model namespace'
-                         ' or the `hparams` namespace/dict.')
+        raise ValueError(f'{attribute} is neither stored in the model namespace'
+                         ' nor the `hparams` namespace/dict, nor the datamodule.')
     return attr
 
 
 def lightning_setattr(model, attribute, value):
     """ Special setattr for lightning. Checks for attribute in model namespace
-        and the old hparams namespace/dict """
+        and the old hparams namespace/dict.
+        Will also set the attribute on datamodule, if it exists.
+    """
+    found = False
     # Check if attribute in model
     if hasattr(model, attribute):
         setattr(model, attribute, value)
+        found = True
     # Check if attribute in model.hparams, either namespace or dict
     elif hasattr(model, 'hparams'):
         if isinstance(model.hparams, dict):
             model.hparams[attribute] = value
         else:
             setattr(model.hparams, attribute, value)
-    else:
-        raise ValueError(f'{attribute} is not stored in the model namespace'
-                         ' or the `hparams` namespace/dict.')
+        found = True
+    # Check if attribute in datamodule
+    if hasattr(model.datamodule, attribute):
+        setattr(model.datamodule, attribute, value)
+        found = True
+
+    if not found:
+        raise ValueError(f'{attribute} is neither stored in the model namespace'
+                         ' nor the `hparams` namespace/dict, nor the datamodule.')
diff --git a/tests/trainer/test_trainer_tricks.py b/tests/trainer/test_trainer_tricks.py
index 66f5b720016cc..ff75563423ea2 100755
--- a/tests/trainer/test_trainer_tricks.py
+++ b/tests/trainer/test_trainer_tricks.py
@@ -8,6 +8,7 @@
 from pytorch_lightning.utilities import AMPType, NATIVE_AMP_AVALAIBLE
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.base import EvalModelTemplate
+from tests.base.datamodules import MNISTDataModule
 
 
 def test_num_training_batches(tmpdir):
@@ -230,11 +231,13 @@ def dataloader(self, *args, **kwargs):
 
     model_class = HparamsEvalModelTemplate if use_hparams else EvalModelTemplate
     model = model_class(**hparams)
+    datamodule = MNISTDataModule(data_dir=tmpdir, batch_size=before_batch_size)
 
     trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, auto_scale_batch_size=True)
-    trainer.fit(model)
+    trainer.fit(model, datamodule)
     after_batch_size = model.hparams.batch_size if use_hparams else model.batch_size
     assert before_batch_size != after_batch_size
+    assert datamodule.batch_size == after_batch_size
 
 
 def test_auto_scale_batch_size_duplicate_attribute_warning(tmpdir):


Code review comments:
diff_hunk: @@ -277,16 +279,12 @@ def _adjust_batch_size(trainer,
     """
     model = trainer.get_model()
     batch_size = lightning_getattr(model, batch_arg_name)
-    if value:
-        lightning_setattr(model, batch_arg_name, value)
-        new_size = value
-        if desc:
-            log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
-    else:
-        new_size = int(batch_size * factor)
-        if desc:
-            log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
-        lightning_setattr(model, batch_arg_name, new_size)
+    new_size = value if value is not None else int(batch_size * factor)
+    if desc:
+        log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
+    lightning_setattr(model, batch_arg_name, new_size)
+    if trainer.datamodule is not None and hasattr(trainer.datamodule, batch_arg_name):
+        setattr(trainer.datamodule, batch_arg_name, new_size)
path: pytorch_lightning/trainer/training_tricks.py
text: @user1:
Is this necessary, should `lightning_setattr` not take care of this? 

@author:
you mean accessing it through model.trainer.datamodule? Ok, I'll try that

diff_hunk: @@ -204,24 +206,36 @@ def lightning_getattr(model, attribute):
             attr = model.hparams[attribute]
         else:
             attr = getattr(model.hparams, attribute)
+    elif hasattr(model.datamodule, attribute):
+        attr = getattr(model.datamodule, attribute)
     else:
-        raise ValueError(f'{attribute} is not stored in the model namespace'
-                         ' or the `hparams` namespace/dict.')
+        raise ValueError(f'{attribute} is neither stored in the model namespace'
+                         ' nor the `hparams` namespace/dict, nor the datamodule.')
     return attr
 
 
 def lightning_setattr(model, attribute, value):
     """ Special setattr for lightning. Checks for attribute in model namespace
-        and the old hparams namespace/dict """
+        and the old hparams namespace/dict.
+        Will also set the attribute on datamodule, if it exists.
+    """
+    found = False
path: pytorch_lightning/utilities/parsing.py
text: @user1:
the same can we check firs if it has the desired attribute?

@author:
@user1 I moved the check to the top, is it better like this?