[refactor] Address Shuhei's comments

nabenabe0928 · ravinkohli · commit df2cdb32dc48 · 2021-12-08T18:30:44.000+01:00
diff --git a/autoPyTorch/pipeline/components/training/trainer/__init__.py b/autoPyTorch/pipeline/components/training/trainer/__init__.py
@@ -413,12 +413,13 @@ def _fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> 'TrainerChoic
 
         if self.choice.use_stochastic_weight_averaging and self.choice.swa_updated:
             # update batch norm statistics
-            swa_utils.update_bn(X['train_data_loader'], self.choice.swa_model.double())
+            swa_utils.update_bn(loader=X['train_data_loader'], model=self.choice.swa_model.double())
+
             # change model
             update_model_state_dict_from_swa(X['network'], self.choice.swa_model.state_dict())
             if self.choice.use_snapshot_ensemble:
                 for model in self.choice.model_snapshots:
-                    swa_utils.update_bn(X['train_data_loader'], model.double())
+                    swa_utils.update_bn(loader=X['train_data_loader'], model=model.double())
 
         # wrap up -- add score if not evaluating every epoch
         if not self.eval_valid_each_epoch(X):
@@ -490,13 +491,10 @@ def early_stop_handler(self, X: Dict[str, Any]) -> bool:
         if self.checkpoint_dir is None:
             self.checkpoint_dir = tempfile.mkdtemp(dir=X['backend'].temporary_directory)
 
+        target_metrics = 'val_loss'
         if X['val_indices'] is None:
-            if X['X_test'] is not None:
-                epochs_since_best = self.run_summary.get_last_epoch() - self.run_summary.get_best_epoch('test_loss')
-            else:
-                epochs_since_best = self.run_summary.get_last_epoch() - self.run_summary.get_best_epoch('train_loss')
-        else:
-            epochs_since_best = self.run_summary.get_last_epoch() - self.run_summary.get_best_epoch()
+            target_metrics = 'test_loss' if X['X_test'] is not None else 'train_loss'
+        epochs_since_best = self.run_summary.get_last_epoch() - self.run_summary.get_best_epoch(target_metrics)
 
         # Save the checkpoint if there is a new best epoch
         best_path = os.path.join(self.checkpoint_dir, 'best.pth')
@@ -626,11 +624,12 @@ def __str__(self) -> str:
     def _get_search_space_updates(self, prefix: Optional[str] = None) -> Dict[str, HyperparameterSearchSpace]:
         """Get the search space updates with the given prefix
 
-        Keyword Arguments:
-            prefix {str} -- Only return search space updates with given prefix (default: {None})
+        Args:
+            prefix (Optional[str]): Only return search space updates with given prefix
 
         Returns:
-            dict -- Mapping of search space updates. Keys don't contain the prefix.
+            Dict[str, HyperparameterSearchSpace]:
+                Mapping of search space updates. Keys don't contain the prefix.
         """
         updates = super()._get_search_space_updates(prefix=prefix)
 
diff --git a/autoPyTorch/pipeline/components/training/trainer/base_trainer.py b/autoPyTorch/pipeline/components/training/trainer/base_trainer.py
@@ -28,7 +28,7 @@
 from autoPyTorch.pipeline.components.training.metrics.metrics import CLASSIFICATION_METRICS, REGRESSION_METRICS
 from autoPyTorch.pipeline.components.training.trainer.utils import Lookahead
 from autoPyTorch.pipeline.components.training.metrics.utils import calculate_score
-from autoPyTorch.pipeline.components.training.trainer.utils import Lookahead, swa_average_function
+from autoPyTorch.pipeline.components.training.trainer.utils import Lookahead, swa_update
 from autoPyTorch.utils.common import FitRequirement, HyperparameterSearchSpace, add_hyperparameter, get_hyperparameter
 from autoPyTorch.utils.implementations import get_loss_weight_strategy
 
@@ -216,7 +216,7 @@ def __init__(self, weighted_loss: bool = False,
                  use_snapshot_ensemble: bool = True,
                  se_lastk: int = 3,
                  use_lookahead_optimizer: bool = True,
-                 random_state: Optional[Union[np.random.RandomState, int]] = None,
+                 random_state: Optional[np.random.RandomState] = None,
                  swa_model: Optional[torch.nn.Module] = None,
                  model_snapshots: Optional[List[torch.nn.Module]] = None,
                  **lookahead_config: Any) -> None:
@@ -277,13 +277,14 @@ def prepare(
 
         # in case we are using swa, maintain an averaged model,
         if self.use_stochastic_weight_averaging:
-            self.swa_model = swa_utils.AveragedModel(self.model, avg_fn=swa_average_function)
+            self.swa_model = swa_utils.AveragedModel(self.model, avg_fn=swa_update)
 
         # in case we are using se or swa, initialise budget_threshold to know when to start swa or se
         self._budget_threshold = 0
         if self.use_stochastic_weight_averaging or self.use_snapshot_ensemble:
-            assert budget_tracker.max_epochs is not None, "Can only use stochastic weight averaging or snapshot " \
-                                                          "ensemble when budget is epochs"
+            if budget_tracker.max_epochs is None:
+                raise ValueError("Budget for stochastic weight averaging or snapshot ensemble must be `epoch`.")
+
             self._budget_threshold = int(0.75 * budget_tracker.max_epochs)
 
         # in case we are using se, initialise list to store model snapshots
@@ -578,7 +579,7 @@ def get_hyperparameter_search_space(
         dataset_properties: Optional[Dict] = None,
         weighted_loss: HyperparameterSearchSpace = HyperparameterSearchSpace(
             hyperparameter="weighted_loss",
-            value_range=[True, False],
+            value_range=(True, False),
             default_value=True),
         la_steps: HyperparameterSearchSpace = HyperparameterSearchSpace(
             hyperparameter="la_steps",
@@ -610,9 +611,7 @@ def get_hyperparameter_search_space(
         cs = ConfigurationSpace()
 
         add_hyperparameter(cs, use_stochastic_weight_averaging, CategoricalHyperparameter)
-        snapshot_ensemble_flag = False
-        if any(use_snapshot_ensemble.value_range):
-            snapshot_ensemble_flag = True
+        snapshot_ensemble_flag = any(use_snapshot_ensemble.value_range)
 
         use_snapshot_ensemble = get_hyperparameter(use_snapshot_ensemble, CategoricalHyperparameter)
         cs.add_hyperparameter(use_snapshot_ensemble)
@@ -623,9 +622,7 @@ def get_hyperparameter_search_space(
             cond = EqualsCondition(se_lastk, use_snapshot_ensemble, True)
             cs.add_condition(cond)
 
-        lookahead_flag = False
-        if any(use_lookahead_optimizer.value_range):
-            lookahead_flag = True
+        lookahead_flag = any(use_lookahead_optimizer.value_range)
 
         use_lookahead_optimizer = get_hyperparameter(use_lookahead_optimizer, CategoricalHyperparameter)
         cs.add_hyperparameter(use_lookahead_optimizer)
diff --git a/autoPyTorch/pipeline/components/training/trainer/cutout_utils.py b/autoPyTorch/pipeline/components/training/trainer/cutout_utils.py
@@ -60,45 +60,45 @@ def criterion_preparation(self, y_a: np.ndarray, y_b: np.ndarray = None, lam: fl
 
     @staticmethod
     def get_hyperparameter_search_space(
-            dataset_properties: Optional[Dict] = None,
-            weighted_loss: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="weighted_loss",
-                value_range=[True, False],
-                default_value=True),
-            la_steps: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="la_steps",
-                value_range=(5, 10),
-                default_value=6,
-                log=False),
-            la_alpha: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="la_alpha",
-                value_range=(0.5, 0.8),
-                default_value=0.6,
-                log=False),
-            use_lookahead_optimizer: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="use_lookahead_optimizer",
-                value_range=(True, False),
-                default_value=True),
-            use_stochastic_weight_averaging: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="use_stochastic_weight_averaging",
-                value_range=(True, False),
-                default_value=True),
-            use_snapshot_ensemble: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="use_snapshot_ensemble",
-                value_range=(True, False),
-                default_value=True),
-            se_lastk: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="se_lastk",
-                value_range=(3,),
-                default_value=3),
-            patch_ratio: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="patch_ratio",
-                value_range=(0, 1),
-                default_value=0.2),
-            cutout_prob: HyperparameterSearchSpace = HyperparameterSearchSpace(
-                hyperparameter="cutout_prob",
-                value_range=(0, 1),
-                default_value=0.2),
+        dataset_properties: Optional[Dict] = None,
+        weighted_loss: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="weighted_loss",
+            value_range=(True, False),
+            default_value=True),
+        la_steps: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="la_steps",
+            value_range=(5, 10),
+            default_value=6,
+            log=False),
+        la_alpha: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="la_alpha",
+            value_range=(0.5, 0.8),
+            default_value=0.6,
+            log=False),
+        use_lookahead_optimizer: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="use_lookahead_optimizer",
+            value_range=(True, False),
+            default_value=True),
+        use_stochastic_weight_averaging: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="use_stochastic_weight_averaging",
+            value_range=(True, False),
+            default_value=True),
+        use_snapshot_ensemble: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="use_snapshot_ensemble",
+            value_range=(True, False),
+            default_value=True),
+        se_lastk: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="se_lastk",
+            value_range=(3,),
+            default_value=3),
+        patch_ratio: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="patch_ratio",
+            value_range=(0, 1),
+            default_value=0.2),
+        cutout_prob: HyperparameterSearchSpace = HyperparameterSearchSpace(
+            hyperparameter="cutout_prob",
+            value_range=(0, 1),
+            default_value=0.2),
     ) -> ConfigurationSpace:
 
         cs = ConfigurationSpace()
diff --git a/autoPyTorch/pipeline/components/training/trainer/mixup_utils.py b/autoPyTorch/pipeline/components/training/trainer/mixup_utils.py
@@ -61,7 +61,7 @@ def get_hyperparameter_search_space(
         dataset_properties: Optional[Dict] = None,
         weighted_loss: HyperparameterSearchSpace = HyperparameterSearchSpace(
             hyperparameter="weighted_loss",
-            value_range=[True, False],
+            value_range=(True, False),
             default_value=True),
         la_steps: HyperparameterSearchSpace = HyperparameterSearchSpace(
             hyperparameter="la_steps",
diff --git a/autoPyTorch/pipeline/components/training/trainer/utils.py b/autoPyTorch/pipeline/components/training/trainer/utils.py
@@ -34,13 +34,16 @@ def update_model_state_dict_from_swa(model: torch.nn.Module, swa_state_dict: Dic
         model_state[name].copy_(param)
 
 
-def swa_average_function(averaged_model_parameter: torch.nn.parameter.Parameter,
-                         model_parameter: torch.nn.parameter.Parameter,
-                         num_averaged: int) -> torch.nn.parameter.Parameter:
+def swa_update(averaged_model_parameter: torch.nn.parameter.Parameter,
+               model_parameter: torch.nn.parameter.Parameter,
+               num_averaged: int) -> torch.nn.parameter.Parameter:
     """
     Pickling the averaged function causes an error because of
     how pytorch initialises the average function.
     Passing this function fixes the issue.
+    The sequential update is performed via:
+        avg[n + 1] = (avg[n] * n + W[n + 1]) / (n + 1)
+
     Args:
         averaged_model_parameter:
         model_parameter:
