temp rm from init

Author: rohitgr7

pytorch_lightning/callbacks/__init__.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from pytorch_lightning.callbacks.base import Callback
-from pytorch_lightning.callbacks.batch_size_finder import BatchSizeFinder
 from pytorch_lightning.callbacks.device_stats_monitor import DeviceStatsMonitor
 from pytorch_lightning.callbacks.early_stopping import EarlyStopping
 from pytorch_lightning.callbacks.finetuning import BackboneFinetuning, BaseFinetuning
@@ -34,7 +33,6 @@
 __all__ = [
     "BackboneFinetuning",
     "BaseFinetuning",
-    "BatchSizeFinder",
     "Callback",
     "DeviceStatsMonitor",
     "EarlyStopping",

pytorch_lightning/callbacks/batch_size_finder.py

@@ -18,7 +18,6 @@
 Finds optimal batch size
 """
 
-import logging
 import os
 import uuid
 from typing import Optional, Tuple
@@ -31,13 +30,12 @@
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities.cloud_io import get_filesystem
 from pytorch_lightning.utilities.data import has_len_all_ranks
+from pytorch_lightning.utilities.distributed import rank_zero_info
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.memory import garbage_collection_cuda, is_oom_error
 from pytorch_lightning.utilities.parsing import lightning_getattr, lightning_hasattr, lightning_setattr
 from pytorch_lightning.utilities.warnings import rank_zero_warn
 
-log = logging.getLogger(__name__)
-
 
 class BatchSizeFinder(Callback):
     def __init__(self, mode: str = "power", steps_per_trial=3, init_val=2, max_trials=25, batch_arg_name="batch_size"):
@@ -124,11 +122,16 @@ def _run_power_scaling(self, trainer, pl_module, new_size):
         """Batch scaling mode where the size is doubled at each iteration until an OOM error is encountered."""
         for _ in range(self.max_trials):
             garbage_collection_cuda()
-            changed = False
 
             try:
                 self._try_loop_run(trainer)
                 new_size, changed = self._adjust_batch_size(trainer, factor=2.0, desc="succeeded")
+
+                if changed:
+                    # Force the dataloaders to reset as the batch size has changed
+                    self._reset_dataloaders(trainer, pl_module)
+                else:
+                    break
             except RuntimeError as exception:
                 if is_oom_error(exception):
                     garbage_collection_cuda()
@@ -137,12 +140,6 @@ def _run_power_scaling(self, trainer, pl_module, new_size):
                 else:
                     raise  # some other error not memory related
 
-            if changed:
-                # Force the train dataloader to reset as the batch size has changed
-                self._reset_dataloaders(trainer, pl_module)
-            else:
-                break
-
         return new_size
 
     def _run_binary_scaling(self, trainer, pl_module, new_size):
@@ -156,7 +153,6 @@ def _run_binary_scaling(self, trainer, pl_module, new_size):
         count = 0
         while True:
             garbage_collection_cuda()
-            trainer.fit_loop.global_step = 0  # reset after each try
             try:
                 # Try fit
                 self._try_loop_run(trainer)
@@ -174,7 +170,7 @@ def _run_binary_scaling(self, trainer, pl_module, new_size):
                     new_size, changed = self._adjust_batch_size(trainer, factor=2.0, desc="succeeded")
 
                 if changed:
-                    # Force the train dataloader to reset as the batch size has changed
+                    # Force the dataloaders to reset as the batch size has changed
                     self._reset_dataloaders(trainer, pl_module)
                 else:
                     break
@@ -287,7 +283,7 @@ def _adjust_batch_size(
                 new batch size
             value: if a value is given, will override the batch size with this value.
                 Note that the value of `factor` will not have an effect in this case
-            desc: either `succeeded` or `failed`. Used purely for logging
+            desc: either ``"succeeded"`` or ``"failed"``. Used purely for logging
 
         Returns:
             The new batch size for the next trial and a bool that signals whether the
@@ -297,7 +293,7 @@ def _adjust_batch_size(
         batch_size = lightning_getattr(model, self.batch_arg_name)
         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}")
+            rank_zero_info(f"Batch size {batch_size} {desc}, trying batch size {new_size}")
 
         # TODO improve this for CombinedLoader
         if trainer.state.fn == TrainerFn.FITTING: