[add] documentation update in base trainer (#468)

Author: theodorju

autoPyTorch/pipeline/components/training/trainer/base_trainer.py

@@ -45,25 +45,53 @@ def __init__(self,
         An object for tracking when to stop the network training.
         It handles epoch based criteria as well as training based criteria.
 
-        It also allows to define a 'epoch_or_time' budget type, which means,
-        the first of them both which is exhausted, is honored
+        It also allows to define a 'epoch_or_time' budget type, which means, the first of them both which is
+        exhausted, is honored
+
+        Args:
+            budget_type (str):
+                Type of budget to be used when fitting the pipeline.
+                Possible values are 'epochs', 'runtime', or 'epoch_or_time'
+            max_epochs (Optional[int], default=None):
+                Maximum number of epochs to train the pipeline for
+            max_runtime (Optional[int], default=None):
+                Maximum number of seconds to train the pipeline for
         """
         self.start_time = time.time()
         self.budget_type = budget_type
         self.max_epochs = max_epochs
         self.max_runtime = max_runtime
 
     def is_max_epoch_reached(self, epoch: int) -> bool:
+        """
+        For budget type 'epoch' or 'epoch_or_time' return True if the maximum number of epochs is reached.
+
+        Args:
+            epoch (int):
+                the current epoch
 
-        # Make None a method to run without this constrain
+        Returns:
+            bool:
+                True if the current epoch is larger than the maximum epochs, False otherwise.
+                Additionally, returns False if the run is without this constraint.
+        """
+        # Make None a method to run without this constraint
         if self.max_epochs is None:
             return False
         if self.budget_type in ['epochs', 'epoch_or_time'] and epoch > self.max_epochs:
             return True
         return False
 
     def is_max_time_reached(self) -> bool:
-        # Make None a method to run without this constrain
+        """
+        For budget type 'runtime' or 'epoch_or_time' return True if the maximum runtime is reached.
+
+        Returns:
+            bool:
+                True if the maximum runtime is reached, False otherwise.
+                Additionally, returns False if the run is without this constraint.
+        """
+        # Make None a method to run without this constraint
         if self.max_runtime is None:
             return False
         elapsed_time = time.time() - self.start_time
@@ -78,14 +106,22 @@ def __init__(
         total_parameter_count: float,
         trainable_parameter_count: float,
         optimize_metric: Optional[str] = None,
-    ):
+    ) -> None:
         """
         A useful object to track performance per epoch.
 
-        It allows to track train, validation and test information not only for
-        debug, but for research purposes (Like understanding overfit).
+        It allows to track train, validation and test information not only for debug, but for research purposes
+        (Like understanding overfit).
 
         It does so by tracking a metric/loss at the end of each epoch.
+
+        Args:
+            total_parameter_count (float):
+                the total number of parameters of the model
+            trainable_parameter_count (float):
+                only the parameters being optimized
+            optimize_metric (Optional[str], default=None):
+                name of the metric that is used to evaluate a pipeline.
         """
         self.performance_tracker: Dict[str, Dict] = {
             'start_time': {},
@@ -121,8 +157,30 @@ def add_performance(self,
                         test_loss: Optional[float] = None,
                         ) -> None:
         """
-        Tracks performance information about the run, useful for
-        plotting individual runs
+        Tracks performance information about the run, useful for plotting individual runs.
+
+        Args:
+            epoch (int):
+                the current epoch
+            start_time (float):
+                timestamp at the beginning of current epoch
+            end_time (float):
+                timestamp when gathering the information after the current epoch
+            train_loss (float):
+                the training loss
+            train_metrics (Dict[str, float]):
+                training scores for each desired metric
+            val_metrics (Dict[str, float]):
+                validation scores for each desired metric
+            test_metrics (Dict[str, float]):
+                test scores for each desired metric
+            val_loss (Optional[float], default=None):
+                the validation loss
+            test_loss (Optional[float], default=None):
+                the test loss
+
+        Returns:
+            None
         """
         self.performance_tracker['train_loss'][epoch] = train_loss
         self.performance_tracker['val_loss'][epoch] = val_loss
@@ -134,6 +192,18 @@ def add_performance(self,
         self.performance_tracker['test_metrics'][epoch] = test_metrics
 
     def get_best_epoch(self, split_type: str = 'val') -> int:
+        """
+        Get the epoch with the best metric.
+
+        Args:
+            split_type (str, default=val):
+                Which split's metric to consider.
+                Possible values are 'train' or 'val
+
+        Returns:
+            int:
+                the epoch with the best metric
+        """
         # If we compute for optimization, prefer the performance
         # metric to the loss
         if self.optimize_metric is not None:
@@ -159,6 +229,13 @@ def get_best_epoch(self, split_type: str = 'val') -> int:
             )) + 1  # Epochs start at 1
 
     def get_last_epoch(self) -> int:
+        """
+        Get the last epoch.
+
+        Returns:
+            int:
+                the last epoch
+        """
         if 'train_loss' not in self.performance_tracker:
             return 0
         else:
@@ -170,7 +247,8 @@ def repr_last_epoch(self) -> str:
         performance
 
         Returns:
-            str: A nice representation of the last epoch
+            str:
+                A nice representation of the last epoch
         """
         last_epoch = len(self.performance_tracker['train_loss'])
         string = "\n"
@@ -202,30 +280,43 @@ def is_empty(self) -> bool:
         Checks if the object is empty or not
 
         Returns:
-            bool
+            bool:
+                True if the object is empty, False otherwise
         """
         # if train_loss is empty, we can be sure that RunSummary is empty.
         return not bool(self.performance_tracker['train_loss'])
 
 
 class BaseTrainerComponent(autoPyTorchTrainingComponent):
     """
-    Base class for training
+    Base class for training.
+
     Args:
-        weighted_loss (int, default=0): In case for classification, whether to weight
-            the loss function according to the distribution of classes in the target
-        use_stochastic_weight_averaging (bool, default=True): whether to use stochastic
-            weight averaging. Stochastic weight averaging is a simple average of
-            multiple points(model parameters) along the trajectory of SGD. SWA
-            has been proposed in
+        weighted_loss (int, default=0):
+            In case for classification, whether to weight the loss function according to the distribution of classes
+            in the target
+        use_stochastic_weight_averaging (bool, default=True):
+            whether to use stochastic weight averaging. Stochastic weight averaging is a simple average of
+            multiple points(model parameters) along the trajectory of SGD. SWA has been proposed in
             [Averaging Weights Leads to Wider Optima and Better Generalization](https://arxiv.org/abs/1803.05407)
-        use_snapshot_ensemble (bool, default=True): whether to use snapshot
-            ensemble
-        se_lastk (int, default=3): Number of snapshots of the network to maintain
-        use_lookahead_optimizer (bool, default=True): whether to use lookahead
-            optimizer
-        random_state:
-        **lookahead_config:
+        use_snapshot_ensemble (bool, default=True):
+            whether to use snapshot ensemble
+        se_lastk (int, default=3):
+            Number of snapshots of the network to maintain
+        use_lookahead_optimizer (bool, default=True):
+            whether to use lookahead optimizer
+        random_state (Optional[np.random.RandomState]):
+            Object that contains a seed and allows for reproducible results
+        swa_model (Optional[torch.nn.Module], default=None):
+            Averaged model used for Stochastic Weight Averaging
+        model_snapshots (Optional[List[torch.nn.Module]], default=None):
+            List of model snapshots in case snapshot ensemble is used
+        **lookahead_config (Any):
+            keyword arguments for the lookahead optimizer including:
+            la_steps (int):
+                number of lookahead steps
+            la_alpha (float):
+                linear interpolation factor. 1.0 recovers the inner optimizer.
     """
     def __init__(self, weighted_loss: int = 0,
                  use_stochastic_weight_averaging: bool = True,
@@ -336,15 +427,21 @@ def prepare(
 
     def on_epoch_start(self, X: Dict[str, Any], epoch: int) -> None:
         """
-        Optional place holder for AutoPytorch Extensions.
+        Optional placeholder for AutoPytorch Extensions.
+        A user can define what happens on every epoch start or every epoch end.
 
-        An user can define what happens on every epoch start or every epoch end.
+        Args:
+            X (Dict[str, Any]):
+                Dictionary with fitted parameters. It is a message passing mechanism, in which during a transform,
+                a components adds relevant information so that further stages can be properly fitted
+            epoch (int):
+                the current epoch
         """
         pass
 
     def _swa_update(self) -> None:
         """
-        perform swa model update
+        Perform Stochastic Weight Averaging model update
         """
         if self.swa_model is None:
             raise ValueError("SWA model cannot be none when stochastic weight averaging is enabled")
@@ -354,6 +451,7 @@ def _swa_update(self) -> None:
     def _se_update(self, epoch: int) -> None:
         """
         Add latest model or swa_model to model snapshot ensemble
+
         Args:
             epoch (int):
                 current epoch
@@ -373,9 +471,16 @@ def _se_update(self, epoch: int) -> None:
 
     def on_epoch_end(self, X: Dict[str, Any], epoch: int) -> bool:
         """
-        Optional place holder for AutoPytorch Extensions.
-        An user can define what happens on every epoch start or every epoch end.
-        If returns True, the training is stopped
+        Optional placeholder for AutoPytorch Extensions.
+        A user can define what happens on every epoch start or every epoch end.
+        If returns True, the training is stopped.
+
+        Args:
+            X (Dict[str, Any]):
+                Dictionary with fitted parameters. It is a message passing mechanism, in which during a transform,
+                a components adds relevant information so that further stages can be properly fitted
+            epoch (int):
+                the current epoch
 
         """
         if X['is_cyclic_scheduler']:
@@ -421,12 +526,18 @@ def train_epoch(self, train_loader: torch.utils.data.DataLoader, epoch: int,
         Train the model for a single epoch.
 
         Args:
-            train_loader (torch.utils.data.DataLoader): generator of features/label
-            epoch (int): The current epoch used solely for tracking purposes
+            train_loader (torch.utils.data.DataLoader):
+                generator of features/label
+            epoch (int):
+                The current epoch used solely for tracking purposes
+            writer (Optional[SummaryWriter]):
+                Object to keep track of the training loss in an event file
 
         Returns:
-            float: training loss
-            Dict[str, float]: scores for each desired metric
+            float:
+                training loss
+            Dict[str, float]:
+                scores for each desired metric
         """
 
         loss_sum = 0.0
@@ -482,12 +593,16 @@ def train_step(self, data: torch.Tensor, targets: torch.Tensor) -> Tuple[float,
         Allows to train 1 step of gradient descent, given a batch of train/labels
 
         Args:
-            data (torch.Tensor): input features to the network
-            targets (torch.Tensor): ground truth to calculate loss
+            data (torch.Tensor):
+                input features to the network
+            targets (torch.Tensor):
+                ground truth to calculate loss
 
         Returns:
-            torch.Tensor: The predictions of the network
-            float: the loss incurred in the prediction
+            torch.Tensor:
+                The predictions of the network
+            float:
+                the loss incurred in the prediction
         """
         # prepare
         data = data.float().to(self.device)
@@ -513,12 +628,18 @@ def evaluate(self, test_loader: torch.utils.data.DataLoader, epoch: int,
         Evaluate the model in both metrics and criterion
 
         Args:
-            test_loader (torch.utils.data.DataLoader): generator of features/label
-            epoch (int): the current epoch for tracking purposes
+            test_loader (torch.utils.data.DataLoader):
+                generator of features/label
+            epoch (int):
+                the current epoch for tracking purposes
+            writer (Optional[SummaryWriter]):
+                Object to keep track of the test loss in an event file
 
         Returns:
-            float: test loss
-            Dict[str, float]: scores for each desired metric
+            float:
+                test loss
+            Dict[str, float]:
+                scores for each desired metric
         """
         self.model.eval()
 
@@ -576,14 +697,15 @@ def get_class_weights(self, criterion: Type[torch.nn.Module], labels: Union[np.n
     def data_preparation(self, X: torch.Tensor, y: torch.Tensor,
                          ) -> Tuple[torch.Tensor, Dict[str, np.ndarray]]:
         """
-        Depending on the trainer choice, data fed to the network might be pre-processed
-        on a different way. That is, in standard training we provide the data to the
-        network as we receive it to the loader. Some regularization techniques, like mixup
-        alter the data.
+        Depending on the trainer choice, data fed to the network might be pre-processed on a different way. That is,
+        in standard training we provide the data to the network as we receive it to the loader. Some regularization
+        techniques, like mixup alter the data.
 
         Args:
-            X (torch.Tensor): The batch training features
-            y (torch.Tensor): The batch training labels
+            X (torch.Tensor):
+                The batch training features
+            y (torch.Tensor):
+                The batch training labels
 
         Returns:
             torch.Tensor: that processes data
@@ -595,16 +717,21 @@ def data_preparation(self, X: torch.Tensor, y: torch.Tensor,
     def criterion_preparation(self, y_a: torch.Tensor, y_b: torch.Tensor = None, lam: float = 1.0
                               ) -> Callable:  # type: ignore
         """
-        Depending on the trainer choice, the criterion is not directly applied to the
-        traditional y_pred/y_ground_truth pairs, but rather it might have a slight transformation.
+        Depending on the trainer choice, the criterion is not directly applied to the traditional
+        y_pred/y_ground_truth pairs, but rather it might have a slight transformation.
         For example, in the case of mixup training, we need to account for the lambda mixup
 
         Args:
-            kwargs (Dict): an expanded dictionary with modifiers to the
-                                  criterion calculation
+            y_a (torch.Tensor):
+                the batch label of the first training example used in trainer
+            y_b (torch.Tensor, default=None):
+                if applicable, the batch label of the second training example used in trainer
+            lam (float):
+                trainer coefficient
 
         Returns:
-            Callable: a lambda function that contains the new criterion calculation recipe
+            Callable:
+                a lambda function that contains the new criterion calculation recipe
         """
         raise NotImplementedError()