Attach data refactor and tuner bugs [4/n]

CHANGELOG.md

@@ -133,6 +133,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `BasePredictionWriter` callback to implement prediction saving ([#7127](https://github.com/PyTorchLightning/pytorch-lightning/pull/7127))
 
 
+- Added `trainer.tune(scale_batch_size_kwargs, lr_find_kwargs)` arguments to configure the tuning algorithms ([#7258](https://github.com/PyTorchLightning/pytorch-lightning/pull/7258))
+
+
 - Added `tpu_distributed` check for TPU Spawn barrier ([#7241](https://github.com/PyTorchLightning/pytorch-lightning/pull/7241))
 
 
@@ -178,6 +181,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Changed default setting for communication of multi-node training using `DDPShardedPlugin` ([#6937](https://github.com/PyTorchLightning/pytorch-lightning/pull/6937))
 
 
+- `trainer.tune()` now returns the tuning result ([#7258](https://github.com/PyTorchLightning/pytorch-lightning/pull/7258))
+
+
 - `LightningModule.from_datasets()` now accepts `IterableDataset` instances as training datasets. ([#7503](https://github.com/PyTorchLightning/pytorch-lightning/pull/7503))
 
 
@@ -325,6 +331,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed torch distributed not available in setup hook for DDP ([#6506](https://github.com/PyTorchLightning/pytorch-lightning/pull/6506))
 
 
+- Fixed `trainer.tuner.{lr_find,scale_batch_size}` not setting the `Trainer` state properly ([#7258](https://github.com/PyTorchLightning/pytorch-lightning/pull/7258))
+
+
 - Fixed bug where `BaseFinetuning.flatten_modules()` was duplicating leaf node parameters ([#6879](https://github.com/PyTorchLightning/pytorch-lightning/pull/6879))
 
 

docs/source/advanced/lr_finder.rst

@@ -73,9 +73,9 @@ If your model is using an arbitrary value instead of ``self.lr`` or ``self.learn
  trainer.tune(model)
 
 
-If you want to inspect the results of the learning rate finder or just play around
-with the parameters of the algorithm, this can be done by invoking the ``lr_find``
-method of the trainer. A typical example of this would look like
+You can also inspect the results of the learning rate finder or just play around
+with the parameters of the algorithm. This can be done by invoking the
+:meth:`~pytorch_lightning.tuner.tuning.Tuner.lr_find` method. A typical example of this would look like:
 
 .. code-block:: python
 

docs/source/advanced/training_tricks.rst

@@ -112,7 +112,7 @@ search for batch sizes larger than the size of the training dataset.
  to `.fit()`.
 
 The scaling algorithm has a number of parameters that the user can control by
-invoking the trainer method `.scale_batch_size` themself (see description below).
+invoking the :meth:`~pytorch_lightning.tuner.tuning.Tuner.scale_batch_size` method:
 
 .. code-block:: python
 
@@ -123,7 +123,7 @@ invoking the trainer method `.scale_batch_size` themself (see description below)
  # Invoke method
  new_batch_size = tuner.scale_batch_size(model, *extra_parameters_here)
 
- # Override old batch size
+ # Override old batch size (this is done automatically)
  model.hparams.batch_size = new_batch_size
 
  # Fit as normal
@@ -142,10 +142,6 @@ The algorithm in short works by:
  3. The found batch size is saved to either `model.batch_size` or `model.hparams.batch_size`
  4. Restore the initial state of model and trainer
 
-.. autoclass:: pytorch_lightning.tuner.tuning.Tuner
- :noindex:
- :members: scale_batch_size
-
 .. warning:: Batch size finder is not supported for DDP yet, it is coming soon.
 
 

docs/source/api_references.rst

@@ -153,14 +153,14 @@ Trainer API
 Tuner API
 ---------
 
-.. currentmodule:: pytorch_lightning.tuner
+.. currentmodule:: pytorch_lightning.tuner.tuning
 
 .. autosummary::
  :toctree: api
  :nosignatures:
+ :template: classtemplate.rst
 
- batch_size_scaling
- lr_finder
+ Tuner
 
 Utilities API
 -------------

pytorch_lightning/core/datamodule.py

@@ -319,7 +319,7 @@ def from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs):
 
  Args:
  args: The parser or namespace to take arguments from. Only known arguments will be
- parsed and passed to the :class:`LightningDataModule`.
+ parsed and passed to the :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
  **kwargs: Additional keyword arguments that may override ones in the parser or namespace.
  These must be valid DataModule arguments.
 

pytorch_lightning/trainer/connectors/data_connector.py

@@ -17,7 +17,6 @@
 from torch.utils.data import DataLoader
 
 import pytorch_lightning as pl
-from pytorch_lightning.core.datamodule import LightningDataModule
 from pytorch_lightning.trainer.supporters import prefetch_iterator
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.model_helpers import is_overridden
@@ -68,24 +67,26 @@ def can_prepare_data(self):
  else:
  return self.trainer.node_rank == 0 and self.trainer.local_rank == 0 and should_call_dm_prepare_data
 
- def attach_data(self, model, train_dataloader, val_dataloaders, datamodule):
- # if a datamodule comes in as the second arg, then fix it for the user
- if isinstance(train_dataloader, LightningDataModule):
- datamodule = train_dataloader
- train_dataloader = None
-
- self.__enforce_datamodule_dataloader_override(train_dataloader, val_dataloaders, datamodule)
-
+ def attach_data(
+ self,
+ model: 'pl.LightningModule',
+ train_dataloader: Optional[Union[DataLoader, List[DataLoader]]] = None,
+ val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,
+ test_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,
+ predict_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,
+ datamodule: Optional['pl.LightningDataModule'] = None
+ ) -> None:
  # set up the passed in dataloaders (if needed)
- self.attach_dataloaders(model, train_dataloader, val_dataloaders)
- self.attach_datamodule(model, datamodule)
-
- def __enforce_datamodule_dataloader_override(self, train_dataloader, val_dataloaders, datamodule):
- # If you supply a datamodule you can't supply train_dataloader or val_dataloaders
- if (train_dataloader is not None or val_dataloaders is not None) and datamodule is not None:
- raise MisconfigurationException(
- 'You cannot pass train_dataloader or val_dataloaders to trainer.fit if you supply a datamodule'
- )
+ self.attach_dataloaders(
+ model,
+ train_dataloader=train_dataloader,
+ val_dataloaders=val_dataloaders,
+ test_dataloaders=test_dataloaders,
+ predict_dataloaders=predict_dataloaders,
+ )
+ self.attach_datamodule(model, datamodule=datamodule)
+ # set local properties on the model
+ self.trainer.model_connector.copy_trainer_model_properties(model)
 
  def attach_dataloaders(
  self,

pytorch_lightning/trainer/trainer.py

@@ -56,8 +56,9 @@
 from pytorch_lightning.trainer.states import TrainerState
 from pytorch_lightning.trainer.training_loop import TrainLoop
 from pytorch_lightning.trainer.training_tricks import TrainerTrainingTricksMixin
+from pytorch_lightning.tuner.lr_finder import _LRFinder
 from pytorch_lightning.tuner.tuning import Tuner
-from pytorch_lightning.utilities import DeviceType, rank_zero_warn
+from pytorch_lightning.utilities import DeviceType, parsing, rank_zero_warn
 from pytorch_lightning.utilities.debugging import InternalDebugger
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.memory import recursive_detach
@@ -409,21 +410,15 @@ def __init__(
  # Callback system
  self.on_init_end()
 
- def _run(
- self,
- model: LightningModule,
- train_dataloader: Any = None,
- val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,
- datamodule: Optional[LightningDataModule] = None,
- ) -> Optional[Union[_EVALUATE_OUTPUT, _PREDICT_OUTPUT]]:
- # set local properties on the model
- self.model_connector.copy_trainer_model_properties(model)
+ def _run(self, model: LightningModule) -> Optional[Union[_EVALUATE_OUTPUT, _PREDICT_OUTPUT]]:
+ # clean hparams
+ if hasattr(model, "hparams"):
+ parsing.clean_namespace(model.hparams)
 
- # ----------------------------
- # LINK DATA
- # ----------------------------
- # setup data, etc...
- self.train_loop.setup_fit(model, train_dataloader, val_dataloaders, datamodule)
+ self.config_validator.verify_loop_configurations(model)
+
+ # attach model log function to callback
+ self.callback_connector.attach_model_logging_functions(model)
 
  # hook
  self.data_connector.prepare_data(model)
@@ -848,14 +843,29 @@ def fit(
  val_dataloaders: Either a single Pytorch Dataloader or a list of them, specifying validation samples.
  If the model has a predefined val_dataloaders method this will be skipped
 
- datamodule: A instance of :class:`LightningDataModule`.
+ datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
  """
  Trainer._log_api_event("fit")
 
  self.state = TrainerState.FITTING
  self.training = True
 
- self._run(model, train_dataloader=train_dataloader, val_dataloaders=val_dataloaders, datamodule=datamodule)
+ # if a datamodule comes in as the second arg, then fix it for the user
+ if isinstance(train_dataloader, LightningDataModule):
+ datamodule = train_dataloader
+ train_dataloader = None
+ # If you supply a datamodule you can't supply train_dataloader or val_dataloaders
+ if (train_dataloader is not None or val_dataloaders is not None) and datamodule is not None:
+ raise MisconfigurationException(
+ 'You cannot pass `train_dataloader` or `val_dataloaders` to `trainer.fit(datamodule=...)`'
+ )
+
+ # links data to the trainer
+ self.data_connector.attach_data(
+ model, train_dataloader=train_dataloader, val_dataloaders=val_dataloaders, datamodule=datamodule
+ )
+
+ self._run(model)
 
  assert self.state.stopped
  self.training = False
@@ -883,7 +893,7 @@ def validate(
 
  verbose: If True, prints the validation results.
 
- datamodule: A instance of :class:`LightningDataModule`.
+ datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
 
  Returns:
  The dictionary with final validation results returned by validation_epoch_end.
@@ -908,10 +918,8 @@ def validate(
  model_provided = model is not None
  model = model or self.lightning_module
 
- # Attach datamodule to get setup/prepare_data added to model before the call to it below
- self.data_connector.attach_datamodule(model, datamodule)
- # Attach dataloaders (if given)
- self.data_connector.attach_dataloaders(model, val_dataloaders=val_dataloaders)
+ # links data to the trainer
+ self.data_connector.attach_data(model, val_dataloaders=val_dataloaders, datamodule=datamodule)
 
  if not model_provided:
  self.validated_ckpt_path = self.__load_ckpt_weights(ckpt_path)
@@ -948,7 +956,7 @@ def test(
 
  verbose: If True, prints the test results.
 
- datamodule: A instance of :class:`LightningDataModule`.
+ datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
 
  Returns:
  Returns a list of dictionaries, one for each test dataloader containing their respective metrics.
@@ -969,10 +977,8 @@ def test(
  model_provided = model is not None
  model = model or self.lightning_module
 
- # Attach datamodule to get setup/prepare_data added to model before the call to it below
- self.data_connector.attach_datamodule(model, datamodule)
- # Attach dataloaders (if given)
- self.data_connector.attach_dataloaders(model, test_dataloaders=test_dataloaders)
+ # links data to the trainer
+ self.data_connector.attach_data(model, test_dataloaders=test_dataloaders, datamodule=datamodule)
 
  if not model_provided:
  self.tested_ckpt_path = self.__load_ckpt_weights(ckpt_path)
@@ -1063,10 +1069,8 @@ def predict(
  if dataloaders is not None and datamodule:
  raise MisconfigurationException('You cannot pass both `trainer.predict(dataloaders=..., datamodule=...)`')
 
- # Attach datamodule to get setup/prepare_data added to model before the call to it below
- self.data_connector.attach_datamodule(model, datamodule)
- # Attach dataloaders (if given)
- self.data_connector.attach_dataloaders(model, predict_dataloaders=dataloaders)
+ # links data to the trainer
+ self.data_connector.attach_data(model, predict_dataloaders=dataloaders, datamodule=datamodule)
 
  results = self._run(model)
 
@@ -1081,7 +1085,9 @@ def tune(
  train_dataloader: Optional[DataLoader] = None,
  val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,
  datamodule: Optional[LightningDataModule] = None,
- ) -> None:
+ scale_batch_size_kwargs: Optional[Dict[str, Any]] = None,
+ lr_find_kwargs: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Optional[Union[int, _LRFinder]]]:
  r"""
  Runs routines to tune hyperparameters before training.
 
@@ -1094,17 +1100,38 @@ def tune(
  val_dataloaders: Either a single Pytorch Dataloader or a list of them, specifying validation samples.
  If the model has a predefined val_dataloaders method this will be skipped
 
- datamodule: A instance of :class:`LightningDataModule`.
+ datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
+
+ scale_batch_size_kwargs: Arguments for :func:`~pytorch_lightning.tuner.batch_size_scaling.scale_batch_size`
+
+ lr_find_kwargs: Arguments for :func:`~pytorch_lightning.tuner.lr_finder.lr_find`
  """
  Trainer._log_api_event("tune")
  self.state = TrainerState.TUNING
  self.tuning = True
 
- self.tuner.tune(model, train_dataloader, val_dataloaders, datamodule)
+ # if a datamodule comes in as the second arg, then fix it for the user
+ if isinstance(train_dataloader, LightningDataModule):
+ datamodule = train_dataloader
+ train_dataloader = None
+ # If you supply a datamodule you can't supply train_dataloader or val_dataloaders
+ if (train_dataloader is not None or val_dataloaders is not None) and datamodule is not None:
+ raise MisconfigurationException(
+ 'You cannot pass `train_dataloader` or `val_dataloaders` to `trainer.tune(datamodule=...)`'
+ )
+
+ # links data to the trainer
+ self.data_connector.attach_data(
+ model, train_dataloader=train_dataloader, val_dataloaders=val_dataloaders, datamodule=datamodule
+ )
+
+ result = self.tuner._tune(model, scale_batch_size_kwargs=scale_batch_size_kwargs, lr_find_kwargs=lr_find_kwargs)
 
  assert self.state.stopped
  self.tuning = False
 
+ return result
+
  def call_setup_hook(self, model: LightningModule) -> None:
  assert self.state.running, f"TrainerState: {self.state}"
  state = self._setup_state

pytorch_lightning/trainer/training_loop.py

@@ -24,7 +24,7 @@
 from pytorch_lightning.plugins import ParallelPlugin
 from pytorch_lightning.trainer.states import TrainerState
 from pytorch_lightning.trainer.supporters import TensorRunningAccum
-from pytorch_lightning.utilities import _TPU_AVAILABLE, AMPType, DeviceType, parsing
+from pytorch_lightning.utilities import _TPU_AVAILABLE, AMPType, DeviceType
 from pytorch_lightning.utilities.distributed import rank_zero_info
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.finite_checks import detect_nan_parameters
@@ -91,20 +91,6 @@ def on_train_start(self):
  # hook
  self.trainer.call_hook("on_train_start")
 
- def setup_fit(self, model, train_dataloader=None, val_dataloaders=None, datamodule=None):
- # clean hparams
- if hasattr(model, "hparams"):
- parsing.clean_namespace(model.hparams)
-
- # links data to the trainer
- self.trainer.data_connector.attach_data(model, train_dataloader, val_dataloaders, datamodule)
-
- # check that model is configured correctly
- self.trainer.config_validator.verify_loop_configurations(model)
-
- # attach model log function to callback
- self.trainer.callback_connector.attach_model_logging_functions(model)
-
  def on_train_end(self):
  if self._teardown_already_run:
  return