We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
auto_scale_batch_size won't reset current_epoch
## 🐛 Bug

When `auto_scale_batch_size` is enabled, the model is initially trained with varying batch sizes. When training begins, `trainer.current_epoch` equals 1 instead of 0.

### To Reproduce

Either observe the progress bar or use a simple callback to track the epoch number, once with  `auto_scale_batch_size`  enabled and once with  `auto_scale_batch_size` disabled.

```
from pytorch_lightning import Callback

class PrintCallback(Callback):
    
    def __init__(self):
        self.observed_epochs = []
        
    def on_train_epoch_start(self, trainer, pl_module):
        print(f'Current Epoch: {trainer.current_epoch}')
        self.observed_epochs.append(trainer.current_epoch)

```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_lightning/tuner/batch_size_scaling.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License
14 import os
15 from pytorch_lightning.core.lightning import LightningModule
16 from pytorch_lightning.utilities.data import has_len
17 from pytorch_lightning.utilities.parsing import lightning_hasattr, lightning_getattr, lightning_setattr
18 from pytorch_lightning.utilities import rank_zero_warn
19 from pytorch_lightning.utilities.exceptions import MisconfigurationException
20 from pytorch_lightning.utilities.memory import is_oom_error, garbage_collection_cuda
21 from pytorch_lightning.loggers.base import DummyLogger
22 from pytorch_lightning import _logger as log
23 from typing import Optional, Tuple
24 
25 
26 def scale_batch_size(trainer,
27                      model: LightningModule,
28                      mode: str = 'power',
29                      steps_per_trial: int = 3,
30                      init_val: int = 2,
31                      max_trials: int = 25,
32                      batch_arg_name: str = 'batch_size',
33                      **fit_kwargs):
34     r"""
35     Will iteratively try to find the largest batch size for a given model
36     that does not give an out of memory (OOM) error.
37 
38     Args:
39         trainer: The Trainer
40         model: Model to fit.
41 
42         mode: string setting the search mode. Either `power` or `binsearch`.
43             If mode is `power` we keep multiplying the batch size by 2, until
44             we get an OOM error. If mode is 'binsearch', we will initially
45             also keep multiplying by 2 and after encountering an OOM error
46             do a binary search between the last successful batch size and the
47             batch size that failed.
48 
49         steps_per_trial: number of steps to run with a given batch size.
50             Idealy 1 should be enough to test if a OOM error occurs,
51             however in practise a few are needed
52 
53         init_val: initial batch size to start the search with
54 
55         max_trials: max number of increase in batch size done before
56            algorithm is terminated
57 
58         batch_arg_name: name of the attribute that stores the batch size.
59             It is expected that the user has provided a model or datamodule that has a hyperparameter
60             with that name. We will look for this attribute name in the following places
61 
62             - `model`
63             - `model.hparams`
64             - `model.datamodule`
65             - `trainer.datamodule` (the datamodule passed to the tune method)
66 
67         **fit_kwargs: remaining arguments to be passed to .fit(), e.g., dataloader
68             or datamodule.
69     """
70     if not lightning_hasattr(model, batch_arg_name):
71         raise MisconfigurationException(
72             f'Field {batch_arg_name} not found in both `model` and `model.hparams`')
73     if hasattr(model, batch_arg_name) and hasattr(model, "hparams") and batch_arg_name in model.hparams:
74         rank_zero_warn(
75             f'Field `model.{batch_arg_name}` and `model.hparams.{batch_arg_name}` are mutually exclusive!'
76             f' `model.{batch_arg_name}` will be used as the initial batch size for scaling.'
77             f' If this is not the intended behavior, please remove either one.'
78         )
79 
80     if hasattr(model.train_dataloader, 'patch_loader_code'):
81         raise MisconfigurationException('The batch scaling feature cannot be used with dataloaders'
82                                         ' passed directly to `.fit()`. Please disable the feature or'
83                                         ' incorporate the dataloader into the model.')
84 
85     # Arguments we adjust during the batch size finder, save for restoring
86     __scale_batch_dump_params(trainer)
87 
88     # Set to values that are required by the algorithm
89     __scale_batch_reset_params(trainer, model, steps_per_trial)
90 
91     # Save initial model, that is loaded after batch size is found
92     save_path = os.path.join(trainer.default_root_dir, 'temp_model.ckpt')
93     trainer.save_checkpoint(str(save_path))
94 
95     if trainer.progress_bar_callback:
96         trainer.progress_bar_callback.disable()
97 
98     # Initially we just double in size until an OOM is encountered
99     new_size = _adjust_batch_size(trainer, value=init_val)  # initially set to init_val
100     if mode == 'power':
101         new_size = _run_power_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs)
102     elif mode == 'binsearch':
103         new_size = _run_binsearch_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs)
104     else:
105         raise ValueError('mode in method `scale_batch_size` can only be `power` or `binsearch')
106 
107     garbage_collection_cuda()
108     log.info(f'Finished batch size finder, will continue with full run using batch size {new_size}')
109 
110     # Restore initial state of model
111     trainer.checkpoint_connector.restore(str(save_path), on_gpu=trainer.on_gpu)
112     os.remove(save_path)
113 
114     # Finish by resetting variables so trainer is ready to fit model
115     __scale_batch_restore_params(trainer)
116     if trainer.progress_bar_callback:
117         trainer.progress_bar_callback.enable()
118 
119     return new_size
120 
121 
122 def __scale_batch_dump_params(trainer):
123     # Prevent going into infinite loop
124     trainer.__dumped_params = {
125         'auto_lr_find': trainer.auto_lr_find,
126         'max_steps': trainer.max_steps,
127         'weights_summary': trainer.weights_summary,
128         'logger': trainer.logger,
129         'callbacks': trainer.callbacks,
130         'checkpoint_callback': trainer.checkpoint_callback,
131         'early_stop_callback': trainer.early_stop_callback,
132         'auto_scale_batch_size': trainer.auto_scale_batch_size,
133         'limit_train_batches': trainer.limit_train_batches,
134         'model': trainer.model,
135     }
136 
137 
138 def __scale_batch_reset_params(trainer, model, steps_per_trial):
139     trainer.auto_scale_batch_size = None  # prevent recursion
140     trainer.auto_lr_find = False  # avoid lr find being called multiple times
141     trainer.max_steps = steps_per_trial  # take few steps
142     trainer.weights_summary = None  # not needed before full run
143     trainer.logger = DummyLogger()
144     trainer.callbacks = []  # not needed before full run
145     trainer.checkpoint_callback = False  # required for saving
146     trainer.early_stop_callback = None
147     trainer.limit_train_batches = 1.0
148     trainer.optimizers, trainer.schedulers = [], []  # required for saving
149     trainer.model = model  # required for saving
150 
151 
152 def __scale_batch_restore_params(trainer):
153     trainer.auto_lr_find = trainer.__dumped_params['auto_lr_find']
154     trainer.max_steps = trainer.__dumped_params['max_steps']
155     trainer.weights_summary = trainer.__dumped_params['weights_summary']
156     trainer.logger = trainer.__dumped_params['logger']
157     trainer.callbacks = trainer.__dumped_params['callbacks']
158     trainer.checkpoint_callback = trainer.__dumped_params['checkpoint_callback']
159     trainer.auto_scale_batch_size = trainer.__dumped_params['auto_scale_batch_size']
160     trainer.early_stop_callback = trainer.__dumped_params['early_stop_callback']
161     trainer.limit_train_batches = trainer.__dumped_params['limit_train_batches']
162     trainer.model = trainer.__dumped_params['model']
163     del trainer.__dumped_params
164 
165 
166 def _run_power_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs):
167     """ Batch scaling mode where the size is doubled at each iteration until an
168         OOM error is encountered. """
169     for _ in range(max_trials):
170         garbage_collection_cuda()
171         trainer.global_step = 0  # reset after each try
172         try:
173             # Try fit
174             trainer.fit(model, **fit_kwargs)
175             # Double in size
176             new_size, changed = _adjust_batch_size(trainer, batch_arg_name, factor=2.0, desc='succeeded')
177         except RuntimeError as exception:
178             # Only these errors should trigger an adjustment
179             if is_oom_error(exception):
180                 # If we fail in power mode, half the size and return
181                 garbage_collection_cuda()
182                 new_size, _ = _adjust_batch_size(trainer, batch_arg_name, factor=0.5, desc='failed')
183                 break
184             else:
185                 raise  # some other error not memory related
186 
187         if not changed:
188             break
189     return new_size
190 
191 
192 def _run_binsearch_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs):
193     """ Batch scaling mode where the size is initially is doubled at each iteration
194         until an OOM error is encountered. Hereafter, the batch size is further
195         refined using a binary search """
196     high = None
197     count = 0
198     while True:
199         garbage_collection_cuda()
200         trainer.global_step = 0  # reset after each try
201         try:
202             # Try fit
203             trainer.fit(model, **fit_kwargs)
204             count += 1
205             if count > max_trials:
206                 break
207             # Double in size
208             low = new_size
209             if high:
210                 if high - low <= 1:
211                     break
212                 midval = (high + low) // 2
213                 new_size, changed = _adjust_batch_size(trainer, batch_arg_name, value=midval, desc='succeeded')
214             else:
215                 new_size, changed = _adjust_batch_size(trainer, batch_arg_name, factor=2.0, desc='succeeded')
216 
217             if not changed:
218                 break
219 
220         except RuntimeError as exception:
221             # Only these errors should trigger an adjustment
222             if is_oom_error(exception):
223                 # If we fail in power mode, half the size and return
224                 garbage_collection_cuda()
225                 high = new_size
226                 midval = (high + low) // 2
227                 new_size, _ = _adjust_batch_size(trainer, value=midval, desc='failed')
228                 if high - low <= 1:
229                     break
230             else:
231                 raise  # some other error not memory related
232 
233     return new_size
234 
235 
236 def _adjust_batch_size(trainer,
237                        batch_arg_name: str = 'batch_size',
238                        factor: float = 1.0,
239                        value: Optional[int] = None,
240                        desc: str = None) -> Tuple[int, bool]:
241     """ Helper function for adjusting the batch size.
242 
243     Args:
244         trainer: instance of pytorch_lightning.Trainer
245 
246         batch_arg_name: name of the field where batch_size is stored.
247 
248         factor: value which the old batch size is multiplied by to get the
249             new batch size
250 
251         value: if a value is given, will override the batch size with this value.
252             Note that the value of `factor` will not have an effect in this case
253 
254         desc: either `succeeded` or `failed`. Used purely for logging
255 
256     Returns:
257         The new batch size for the next trial and a bool that signals whether the
258         new value is different than the previous batch size.
259     """
260     model = trainer.get_model()
261     batch_size = lightning_getattr(model, batch_arg_name)
262     new_size = value if value is not None else int(batch_size * factor)
263     if desc:
264         log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')
265 
266     if not _is_valid_batch_size(new_size, trainer.train_dataloader):
267         new_size = min(new_size, len(trainer.train_dataloader.dataset))
268 
269     changed = new_size != batch_size
270     lightning_setattr(model, batch_arg_name, new_size)
271     return new_size, changed
272 
273 
274 def _is_valid_batch_size(current_size, dataloader):
275     return not has_len(dataloader) or current_size <= len(dataloader)
276
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pytorch_lightning/tuner/batch_size_scaling.py b/pytorch_lightning/tuner/batch_size_scaling.py
--- a/pytorch_lightning/tuner/batch_size_scaling.py
+++ b/pytorch_lightning/tuner/batch_size_scaling.py
@@ -123,6 +123,7 @@
     # Prevent going into infinite loop
     trainer.__dumped_params = {
         'auto_lr_find': trainer.auto_lr_find,
+        'current_epoch': trainer.current_epoch,
         'max_steps': trainer.max_steps,
         'weights_summary': trainer.weights_summary,
         'logger': trainer.logger,
@@ -138,6 +139,7 @@
 def __scale_batch_reset_params(trainer, model, steps_per_trial):
     trainer.auto_scale_batch_size = None  # prevent recursion
     trainer.auto_lr_find = False  # avoid lr find being called multiple times
+    trainer.current_epoch = 0
     trainer.max_steps = steps_per_trial  # take few steps
     trainer.weights_summary = None  # not needed before full run
     trainer.logger = DummyLogger()
@@ -151,6 +153,7 @@
 
 def __scale_batch_restore_params(trainer):
     trainer.auto_lr_find = trainer.__dumped_params['auto_lr_find']
+    trainer.current_epoch = trainer.__dumped_params['current_epoch']
     trainer.max_steps = trainer.__dumped_params['max_steps']
     trainer.weights_summary = trainer.__dumped_params['weights_summary']
     trainer.logger = trainer.__dumped_params['logger']

{"golden_diff": "diff --git a/pytorch_lightning/tuner/batch_size_scaling.py b/pytorch_lightning/tuner/batch_size_scaling.py\n--- a/pytorch_lightning/tuner/batch_size_scaling.py\n+++ b/pytorch_lightning/tuner/batch_size_scaling.py\n@@ -123,6 +123,7 @@\n     # Prevent going into infinite loop\n     trainer.__dumped_params = {\n         'auto_lr_find': trainer.auto_lr_find,\n+        'current_epoch': trainer.current_epoch,\n         'max_steps': trainer.max_steps,\n         'weights_summary': trainer.weights_summary,\n         'logger': trainer.logger,\n@@ -138,6 +139,7 @@\n def __scale_batch_reset_params(trainer, model, steps_per_trial):\n     trainer.auto_scale_batch_size = None  # prevent recursion\n     trainer.auto_lr_find = False  # avoid lr find being called multiple times\n+    trainer.current_epoch = 0\n     trainer.max_steps = steps_per_trial  # take few steps\n     trainer.weights_summary = None  # not needed before full run\n     trainer.logger = DummyLogger()\n@@ -151,6 +153,7 @@\n \n def __scale_batch_restore_params(trainer):\n     trainer.auto_lr_find = trainer.__dumped_params['auto_lr_find']\n+    trainer.current_epoch = trainer.__dumped_params['current_epoch']\n     trainer.max_steps = trainer.__dumped_params['max_steps']\n     trainer.weights_summary = trainer.__dumped_params['weights_summary']\n     trainer.logger = trainer.__dumped_params['logger']\n", "issue": "auto_scale_batch_size won't reset current_epoch\n## \ud83d\udc1b Bug\r\n\r\nWhen `auto_scale_batch_size` is enabled, the model is initially trained with varying batch sizes. When training begins, `trainer.current_epoch` equals 1 instead of 0.\r\n\r\n### To Reproduce\r\n\r\nEither observe the progress bar or use a simple callback to track the epoch number, once with  `auto_scale_batch_size`  enabled and once with  `auto_scale_batch_size` disabled.\r\n\r\n```\r\nfrom pytorch_lightning import Callback\r\n\r\nclass PrintCallback(Callback):\r\n    \r\n    def __init__(self):\r\n        self.observed_epochs = []\r\n        \r\n    def on_train_epoch_start(self, trainer, pl_module):\r\n        print(f'Current Epoch: {trainer.current_epoch}')\r\n        self.observed_epochs.append(trainer.current_epoch)\r\n\r\n```\n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License\nimport os\nfrom pytorch_lightning.core.lightning import LightningModule\nfrom pytorch_lightning.utilities.data import has_len\nfrom pytorch_lightning.utilities.parsing import lightning_hasattr, lightning_getattr, lightning_setattr\nfrom pytorch_lightning.utilities import rank_zero_warn\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom pytorch_lightning.utilities.memory import is_oom_error, garbage_collection_cuda\nfrom pytorch_lightning.loggers.base import DummyLogger\nfrom pytorch_lightning import _logger as log\nfrom typing import Optional, Tuple\n\n\ndef scale_batch_size(trainer,\n                     model: LightningModule,\n                     mode: str = 'power',\n                     steps_per_trial: int = 3,\n                     init_val: int = 2,\n                     max_trials: int = 25,\n                     batch_arg_name: str = 'batch_size',\n                     **fit_kwargs):\n    r\"\"\"\n    Will iteratively try to find the largest batch size for a given model\n    that does not give an out of memory (OOM) error.\n\n    Args:\n        trainer: The Trainer\n        model: Model to fit.\n\n        mode: string setting the search mode. Either `power` or `binsearch`.\n            If mode is `power` we keep multiplying the batch size by 2, until\n            we get an OOM error. If mode is 'binsearch', we will initially\n            also keep multiplying by 2 and after encountering an OOM error\n            do a binary search between the last successful batch size and the\n            batch size that failed.\n\n        steps_per_trial: number of steps to run with a given batch size.\n            Idealy 1 should be enough to test if a OOM error occurs,\n            however in practise a few are needed\n\n        init_val: initial batch size to start the search with\n\n        max_trials: max number of increase in batch size done before\n           algorithm is terminated\n\n        batch_arg_name: name of the attribute that stores the batch size.\n            It is expected that the user has provided a model or datamodule that has a hyperparameter\n            with that name. We will look for this attribute name in the following places\n\n            - `model`\n            - `model.hparams`\n            - `model.datamodule`\n            - `trainer.datamodule` (the datamodule passed to the tune method)\n\n        **fit_kwargs: remaining arguments to be passed to .fit(), e.g., dataloader\n            or datamodule.\n    \"\"\"\n    if not lightning_hasattr(model, batch_arg_name):\n        raise MisconfigurationException(\n            f'Field {batch_arg_name} not found in both `model` and `model.hparams`')\n    if hasattr(model, batch_arg_name) and hasattr(model, \"hparams\") and batch_arg_name in model.hparams:\n        rank_zero_warn(\n            f'Field `model.{batch_arg_name}` and `model.hparams.{batch_arg_name}` are mutually exclusive!'\n            f' `model.{batch_arg_name}` will be used as the initial batch size for scaling.'\n            f' If this is not the intended behavior, please remove either one.'\n        )\n\n    if hasattr(model.train_dataloader, 'patch_loader_code'):\n        raise MisconfigurationException('The batch scaling feature cannot be used with dataloaders'\n                                        ' passed directly to `.fit()`. Please disable the feature or'\n                                        ' incorporate the dataloader into the model.')\n\n    # Arguments we adjust during the batch size finder, save for restoring\n    __scale_batch_dump_params(trainer)\n\n    # Set to values that are required by the algorithm\n    __scale_batch_reset_params(trainer, model, steps_per_trial)\n\n    # Save initial model, that is loaded after batch size is found\n    save_path = os.path.join(trainer.default_root_dir, 'temp_model.ckpt')\n    trainer.save_checkpoint(str(save_path))\n\n    if trainer.progress_bar_callback:\n        trainer.progress_bar_callback.disable()\n\n    # Initially we just double in size until an OOM is encountered\n    new_size = _adjust_batch_size(trainer, value=init_val)  # initially set to init_val\n    if mode == 'power':\n        new_size = _run_power_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs)\n    elif mode == 'binsearch':\n        new_size = _run_binsearch_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs)\n    else:\n        raise ValueError('mode in method `scale_batch_size` can only be `power` or `binsearch')\n\n    garbage_collection_cuda()\n    log.info(f'Finished batch size finder, will continue with full run using batch size {new_size}')\n\n    # Restore initial state of model\n    trainer.checkpoint_connector.restore(str(save_path), on_gpu=trainer.on_gpu)\n    os.remove(save_path)\n\n    # Finish by resetting variables so trainer is ready to fit model\n    __scale_batch_restore_params(trainer)\n    if trainer.progress_bar_callback:\n        trainer.progress_bar_callback.enable()\n\n    return new_size\n\n\ndef __scale_batch_dump_params(trainer):\n    # Prevent going into infinite loop\n    trainer.__dumped_params = {\n        'auto_lr_find': trainer.auto_lr_find,\n        'max_steps': trainer.max_steps,\n        'weights_summary': trainer.weights_summary,\n        'logger': trainer.logger,\n        'callbacks': trainer.callbacks,\n        'checkpoint_callback': trainer.checkpoint_callback,\n        'early_stop_callback': trainer.early_stop_callback,\n        'auto_scale_batch_size': trainer.auto_scale_batch_size,\n        'limit_train_batches': trainer.limit_train_batches,\n        'model': trainer.model,\n    }\n\n\ndef __scale_batch_reset_params(trainer, model, steps_per_trial):\n    trainer.auto_scale_batch_size = None  # prevent recursion\n    trainer.auto_lr_find = False  # avoid lr find being called multiple times\n    trainer.max_steps = steps_per_trial  # take few steps\n    trainer.weights_summary = None  # not needed before full run\n    trainer.logger = DummyLogger()\n    trainer.callbacks = []  # not needed before full run\n    trainer.checkpoint_callback = False  # required for saving\n    trainer.early_stop_callback = None\n    trainer.limit_train_batches = 1.0\n    trainer.optimizers, trainer.schedulers = [], []  # required for saving\n    trainer.model = model  # required for saving\n\n\ndef __scale_batch_restore_params(trainer):\n    trainer.auto_lr_find = trainer.__dumped_params['auto_lr_find']\n    trainer.max_steps = trainer.__dumped_params['max_steps']\n    trainer.weights_summary = trainer.__dumped_params['weights_summary']\n    trainer.logger = trainer.__dumped_params['logger']\n    trainer.callbacks = trainer.__dumped_params['callbacks']\n    trainer.checkpoint_callback = trainer.__dumped_params['checkpoint_callback']\n    trainer.auto_scale_batch_size = trainer.__dumped_params['auto_scale_batch_size']\n    trainer.early_stop_callback = trainer.__dumped_params['early_stop_callback']\n    trainer.limit_train_batches = trainer.__dumped_params['limit_train_batches']\n    trainer.model = trainer.__dumped_params['model']\n    del trainer.__dumped_params\n\n\ndef _run_power_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs):\n    \"\"\" Batch scaling mode where the size is doubled at each iteration until an\n        OOM error is encountered. \"\"\"\n    for _ in range(max_trials):\n        garbage_collection_cuda()\n        trainer.global_step = 0  # reset after each try\n        try:\n            # Try fit\n            trainer.fit(model, **fit_kwargs)\n            # Double in size\n            new_size, changed = _adjust_batch_size(trainer, batch_arg_name, factor=2.0, desc='succeeded')\n        except RuntimeError as exception:\n            # Only these errors should trigger an adjustment\n            if is_oom_error(exception):\n                # If we fail in power mode, half the size and return\n                garbage_collection_cuda()\n                new_size, _ = _adjust_batch_size(trainer, batch_arg_name, factor=0.5, desc='failed')\n                break\n            else:\n                raise  # some other error not memory related\n\n        if not changed:\n            break\n    return new_size\n\n\ndef _run_binsearch_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs):\n    \"\"\" Batch scaling mode where the size is initially is doubled at each iteration\n        until an OOM error is encountered. Hereafter, the batch size is further\n        refined using a binary search \"\"\"\n    high = None\n    count = 0\n    while True:\n        garbage_collection_cuda()\n        trainer.global_step = 0  # reset after each try\n        try:\n            # Try fit\n            trainer.fit(model, **fit_kwargs)\n            count += 1\n            if count > max_trials:\n                break\n            # Double in size\n            low = new_size\n            if high:\n                if high - low <= 1:\n                    break\n                midval = (high + low) // 2\n                new_size, changed = _adjust_batch_size(trainer, batch_arg_name, value=midval, desc='succeeded')\n            else:\n                new_size, changed = _adjust_batch_size(trainer, batch_arg_name, factor=2.0, desc='succeeded')\n\n            if not changed:\n                break\n\n        except RuntimeError as exception:\n            # Only these errors should trigger an adjustment\n            if is_oom_error(exception):\n                # If we fail in power mode, half the size and return\n                garbage_collection_cuda()\n                high = new_size\n                midval = (high + low) // 2\n                new_size, _ = _adjust_batch_size(trainer, value=midval, desc='failed')\n                if high - low <= 1:\n                    break\n            else:\n                raise  # some other error not memory related\n\n    return new_size\n\n\ndef _adjust_batch_size(trainer,\n                       batch_arg_name: str = 'batch_size',\n                       factor: float = 1.0,\n                       value: Optional[int] = None,\n                       desc: str = None) -> Tuple[int, bool]:\n    \"\"\" Helper function for adjusting the batch size.\n\n    Args:\n        trainer: instance of pytorch_lightning.Trainer\n\n        batch_arg_name: name of the field where batch_size is stored.\n\n        factor: value which the old batch size is multiplied by to get the\n            new batch size\n\n        value: if a value is given, will override the batch size with this value.\n            Note that the value of `factor` will not have an effect in this case\n\n        desc: either `succeeded` or `failed`. Used purely for logging\n\n    Returns:\n        The new batch size for the next trial and a bool that signals whether the\n        new value is different than the previous batch size.\n    \"\"\"\n    model = trainer.get_model()\n    batch_size = lightning_getattr(model, batch_arg_name)\n    new_size = value if value is not None else int(batch_size * factor)\n    if desc:\n        log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')\n\n    if not _is_valid_batch_size(new_size, trainer.train_dataloader):\n        new_size = min(new_size, len(trainer.train_dataloader.dataset))\n\n    changed = new_size != batch_size\n    lightning_setattr(model, batch_arg_name, new_size)\n    return new_size, changed\n\n\ndef _is_valid_batch_size(current_size, dataloader):\n    return not has_len(dataloader) or current_size <= len(dataloader)\n", "path": "pytorch_lightning/tuner/batch_size_scaling.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License\nimport os\nfrom pytorch_lightning.core.lightning import LightningModule\nfrom pytorch_lightning.utilities.data import has_len\nfrom pytorch_lightning.utilities.parsing import lightning_hasattr, lightning_getattr, lightning_setattr\nfrom pytorch_lightning.utilities import rank_zero_warn\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom pytorch_lightning.utilities.memory import is_oom_error, garbage_collection_cuda\nfrom pytorch_lightning.loggers.base import DummyLogger\nfrom pytorch_lightning import _logger as log\nfrom typing import Optional, Tuple\n\n\ndef scale_batch_size(trainer,\n                     model: LightningModule,\n                     mode: str = 'power',\n                     steps_per_trial: int = 3,\n                     init_val: int = 2,\n                     max_trials: int = 25,\n                     batch_arg_name: str = 'batch_size',\n                     **fit_kwargs):\n    r\"\"\"\n    Will iteratively try to find the largest batch size for a given model\n    that does not give an out of memory (OOM) error.\n\n    Args:\n        trainer: The Trainer\n        model: Model to fit.\n\n        mode: string setting the search mode. Either `power` or `binsearch`.\n            If mode is `power` we keep multiplying the batch size by 2, until\n            we get an OOM error. If mode is 'binsearch', we will initially\n            also keep multiplying by 2 and after encountering an OOM error\n            do a binary search between the last successful batch size and the\n            batch size that failed.\n\n        steps_per_trial: number of steps to run with a given batch size.\n            Idealy 1 should be enough to test if a OOM error occurs,\n            however in practise a few are needed\n\n        init_val: initial batch size to start the search with\n\n        max_trials: max number of increase in batch size done before\n           algorithm is terminated\n\n        batch_arg_name: name of the attribute that stores the batch size.\n            It is expected that the user has provided a model or datamodule that has a hyperparameter\n            with that name. We will look for this attribute name in the following places\n\n            - `model`\n            - `model.hparams`\n            - `model.datamodule`\n            - `trainer.datamodule` (the datamodule passed to the tune method)\n\n        **fit_kwargs: remaining arguments to be passed to .fit(), e.g., dataloader\n            or datamodule.\n    \"\"\"\n    if not lightning_hasattr(model, batch_arg_name):\n        raise MisconfigurationException(\n            f'Field {batch_arg_name} not found in both `model` and `model.hparams`')\n    if hasattr(model, batch_arg_name) and hasattr(model, \"hparams\") and batch_arg_name in model.hparams:\n        rank_zero_warn(\n            f'Field `model.{batch_arg_name}` and `model.hparams.{batch_arg_name}` are mutually exclusive!'\n            f' `model.{batch_arg_name}` will be used as the initial batch size for scaling.'\n            f' If this is not the intended behavior, please remove either one.'\n        )\n\n    if hasattr(model.train_dataloader, 'patch_loader_code'):\n        raise MisconfigurationException('The batch scaling feature cannot be used with dataloaders'\n                                        ' passed directly to `.fit()`. Please disable the feature or'\n                                        ' incorporate the dataloader into the model.')\n\n    # Arguments we adjust during the batch size finder, save for restoring\n    __scale_batch_dump_params(trainer)\n\n    # Set to values that are required by the algorithm\n    __scale_batch_reset_params(trainer, model, steps_per_trial)\n\n    # Save initial model, that is loaded after batch size is found\n    save_path = os.path.join(trainer.default_root_dir, 'temp_model.ckpt')\n    trainer.save_checkpoint(str(save_path))\n\n    if trainer.progress_bar_callback:\n        trainer.progress_bar_callback.disable()\n\n    # Initially we just double in size until an OOM is encountered\n    new_size = _adjust_batch_size(trainer, value=init_val)  # initially set to init_val\n    if mode == 'power':\n        new_size = _run_power_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs)\n    elif mode == 'binsearch':\n        new_size = _run_binsearch_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs)\n    else:\n        raise ValueError('mode in method `scale_batch_size` can only be `power` or `binsearch')\n\n    garbage_collection_cuda()\n    log.info(f'Finished batch size finder, will continue with full run using batch size {new_size}')\n\n    # Restore initial state of model\n    trainer.checkpoint_connector.restore(str(save_path), on_gpu=trainer.on_gpu)\n    os.remove(save_path)\n\n    # Finish by resetting variables so trainer is ready to fit model\n    __scale_batch_restore_params(trainer)\n    if trainer.progress_bar_callback:\n        trainer.progress_bar_callback.enable()\n\n    return new_size\n\n\ndef __scale_batch_dump_params(trainer):\n    # Prevent going into infinite loop\n    trainer.__dumped_params = {\n        'auto_lr_find': trainer.auto_lr_find,\n        'current_epoch': trainer.current_epoch,\n        'max_steps': trainer.max_steps,\n        'weights_summary': trainer.weights_summary,\n        'logger': trainer.logger,\n        'callbacks': trainer.callbacks,\n        'checkpoint_callback': trainer.checkpoint_callback,\n        'early_stop_callback': trainer.early_stop_callback,\n        'auto_scale_batch_size': trainer.auto_scale_batch_size,\n        'limit_train_batches': trainer.limit_train_batches,\n        'model': trainer.model,\n    }\n\n\ndef __scale_batch_reset_params(trainer, model, steps_per_trial):\n    trainer.auto_scale_batch_size = None  # prevent recursion\n    trainer.auto_lr_find = False  # avoid lr find being called multiple times\n    trainer.current_epoch = 0\n    trainer.max_steps = steps_per_trial  # take few steps\n    trainer.weights_summary = None  # not needed before full run\n    trainer.logger = DummyLogger()\n    trainer.callbacks = []  # not needed before full run\n    trainer.checkpoint_callback = False  # required for saving\n    trainer.early_stop_callback = None\n    trainer.limit_train_batches = 1.0\n    trainer.optimizers, trainer.schedulers = [], []  # required for saving\n    trainer.model = model  # required for saving\n\n\ndef __scale_batch_restore_params(trainer):\n    trainer.auto_lr_find = trainer.__dumped_params['auto_lr_find']\n    trainer.current_epoch = trainer.__dumped_params['current_epoch']\n    trainer.max_steps = trainer.__dumped_params['max_steps']\n    trainer.weights_summary = trainer.__dumped_params['weights_summary']\n    trainer.logger = trainer.__dumped_params['logger']\n    trainer.callbacks = trainer.__dumped_params['callbacks']\n    trainer.checkpoint_callback = trainer.__dumped_params['checkpoint_callback']\n    trainer.auto_scale_batch_size = trainer.__dumped_params['auto_scale_batch_size']\n    trainer.early_stop_callback = trainer.__dumped_params['early_stop_callback']\n    trainer.limit_train_batches = trainer.__dumped_params['limit_train_batches']\n    trainer.model = trainer.__dumped_params['model']\n    del trainer.__dumped_params\n\n\ndef _run_power_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs):\n    \"\"\" Batch scaling mode where the size is doubled at each iteration until an\n        OOM error is encountered. \"\"\"\n    for _ in range(max_trials):\n        garbage_collection_cuda()\n        trainer.global_step = 0  # reset after each try\n        try:\n            # Try fit\n            trainer.fit(model, **fit_kwargs)\n            # Double in size\n            new_size, changed = _adjust_batch_size(trainer, batch_arg_name, factor=2.0, desc='succeeded')\n        except RuntimeError as exception:\n            # Only these errors should trigger an adjustment\n            if is_oom_error(exception):\n                # If we fail in power mode, half the size and return\n                garbage_collection_cuda()\n                new_size, _ = _adjust_batch_size(trainer, batch_arg_name, factor=0.5, desc='failed')\n                break\n            else:\n                raise  # some other error not memory related\n\n        if not changed:\n            break\n    return new_size\n\n\ndef _run_binsearch_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs):\n    \"\"\" Batch scaling mode where the size is initially is doubled at each iteration\n        until an OOM error is encountered. Hereafter, the batch size is further\n        refined using a binary search \"\"\"\n    high = None\n    count = 0\n    while True:\n        garbage_collection_cuda()\n        trainer.global_step = 0  # reset after each try\n        try:\n            # Try fit\n            trainer.fit(model, **fit_kwargs)\n            count += 1\n            if count > max_trials:\n                break\n            # Double in size\n            low = new_size\n            if high:\n                if high - low <= 1:\n                    break\n                midval = (high + low) // 2\n                new_size, changed = _adjust_batch_size(trainer, batch_arg_name, value=midval, desc='succeeded')\n            else:\n                new_size, changed = _adjust_batch_size(trainer, batch_arg_name, factor=2.0, desc='succeeded')\n\n            if not changed:\n                break\n\n        except RuntimeError as exception:\n            # Only these errors should trigger an adjustment\n            if is_oom_error(exception):\n                # If we fail in power mode, half the size and return\n                garbage_collection_cuda()\n                high = new_size\n                midval = (high + low) // 2\n                new_size, _ = _adjust_batch_size(trainer, value=midval, desc='failed')\n                if high - low <= 1:\n                    break\n            else:\n                raise  # some other error not memory related\n\n    return new_size\n\n\ndef _adjust_batch_size(trainer,\n                       batch_arg_name: str = 'batch_size',\n                       factor: float = 1.0,\n                       value: Optional[int] = None,\n                       desc: str = None) -> Tuple[int, bool]:\n    \"\"\" Helper function for adjusting the batch size.\n\n    Args:\n        trainer: instance of pytorch_lightning.Trainer\n\n        batch_arg_name: name of the field where batch_size is stored.\n\n        factor: value which the old batch size is multiplied by to get the\n            new batch size\n\n        value: if a value is given, will override the batch size with this value.\n            Note that the value of `factor` will not have an effect in this case\n\n        desc: either `succeeded` or `failed`. Used purely for logging\n\n    Returns:\n        The new batch size for the next trial and a bool that signals whether the\n        new value is different than the previous batch size.\n    \"\"\"\n    model = trainer.get_model()\n    batch_size = lightning_getattr(model, batch_arg_name)\n    new_size = value if value is not None else int(batch_size * factor)\n    if desc:\n        log.info(f'Batch size {batch_size} {desc}, trying batch size {new_size}')\n\n    if not _is_valid_batch_size(new_size, trainer.train_dataloader):\n        new_size = min(new_size, len(trainer.train_dataloader.dataset))\n\n    changed = new_size != batch_size\n    lightning_setattr(model, batch_arg_name, new_size)\n    return new_size, changed\n\n\ndef _is_valid_batch_size(current_size, dataloader):\n    return not has_len(dataloader) or current_size <= len(dataloader)\n", "path": "pytorch_lightning/tuner/batch_size_scaling.py"}]}