Enable learned embeddings, fix bug with non cyclic schedulers

Author: ravinkohli

autoPyTorch/pipeline/components/setup/network_embedding/__init__.py

@@ -148,71 +148,62 @@ def get_hyperparameter_search_space(
         if default is None:
             defaults = [
                 'NoEmbedding',
-                # 'LearnedEntityEmbedding',
+                'LearnedEntityEmbedding',
             ]
             for default_ in defaults:
                 if default_ in available_embedding:
                     default = default_
                     break
 
-        # Restrict embedding to NoEmbedding until preprocessing is fixed
-        embedding = CSH.CategoricalHyperparameter('__choice__',
-                                                  ['NoEmbedding'],
-                                                  default_value=default)
+        categorical_columns = dataset_properties['categorical_columns'] \
+            if isinstance(dataset_properties['categorical_columns'], List) else []
+
+        updates = self._get_search_space_updates()
+        if '__choice__' in updates.keys():
+            choice_hyperparameter = updates['__choice__']
+            if not set(choice_hyperparameter.value_range).issubset(available_embedding):
+                raise ValueError("Expected given update for {} to have "
+                                 "choices in {} got {}".format(self.__class__.__name__,
+                                                               available_embedding,
+                                                               choice_hyperparameter.value_range))
+            if len(categorical_columns) == 0:
+                assert len(choice_hyperparameter.value_range) == 1
+                if 'NoEmbedding' not in choice_hyperparameter.value_range:
+                    raise ValueError("Provided {} in choices, however, the dataset "
+                                     "is incompatible with it".format(choice_hyperparameter.value_range))
+            embedding = CSH.CategoricalHyperparameter('__choice__',
+                                                      choice_hyperparameter.value_range,
+                                                      default_value=choice_hyperparameter.default_value)
+        else:
+
+            if len(categorical_columns) == 0:
+                default = 'NoEmbedding'
+                if include is not None and default not in include:
+                    raise ValueError("Provided {} in include, however, the dataset "
+                                     "is incompatible with it".format(include))
+                embedding = CSH.CategoricalHyperparameter('__choice__',
+                                                          ['NoEmbedding'],
+                                                          default_value=default)
+            else:
+                embedding = CSH.CategoricalHyperparameter('__choice__',
+                                                          list(available_embedding.keys()),
+                                                          default_value=default)
+
         cs.add_hyperparameter(embedding)
+        for name in embedding.choices:
+            updates = self._get_search_space_updates(prefix=name)
+            config_space = available_embedding[name].get_hyperparameter_search_space(dataset_properties,  # type: ignore
+                                                                                     **updates)
+            parent_hyperparameter = {'parent': embedding, 'value': name}
+            cs.add_configuration_space(
+                name,
+                config_space,
+                parent_hyperparameter=parent_hyperparameter
+            )
+
         self.configuration_space_ = cs
         self.dataset_properties_ = dataset_properties
         return cs
-        # categorical_columns = dataset_properties['categorical_columns'] \
-        #     if isinstance(dataset_properties['categorical_columns'], List) else []
-
-        # updates = self._get_search_space_updates()
-        # if '__choice__' in updates.keys():
-        #     choice_hyperparameter = updates['__choice__']
-        #     if not set(choice_hyperparameter.value_range).issubset(available_embedding):
-        #         raise ValueError("Expected given update for {} to have "
-        #                          "choices in {} got {}".format(self.__class__.__name__,
-        #                                                        available_embedding,
-        #                                                        choice_hyperparameter.value_range))
-        #     if len(categorical_columns) == 0:
-        #         assert len(choice_hyperparameter.value_range) == 1
-        #         if 'NoEmbedding' not in choice_hyperparameter.value_range:
-        #             raise ValueError("Provided {} in choices, however, the dataset "
-        #                              "is incompatible with it".format(choice_hyperparameter.value_range))
-        #     embedding = CSH.CategoricalHyperparameter('__choice__',
-        #                                               choice_hyperparameter.value_range,
-        #                                               default_value=choice_hyperparameter.default_value)
-        # else:
-
-        #     if len(categorical_columns) == 0:
-        #         default = 'NoEmbedding'
-        #         if include is not None and default not in include:
-        #             raise ValueError("Provided {} in include, however, the dataset "
-        #                              "is incompatible with it".format(include))
-        #         embedding = CSH.CategoricalHyperparameter('__choice__',
-        #                                                   ['NoEmbedding'],
-        #                                                   default_value=default)
-        #     else:
-        #         embedding = CSH.CategoricalHyperparameter('__choice__',
-        #                                                   list(available_embedding.keys()),
-        #                                                   default_value=default)
-
-        # cs.add_hyperparameter(embedding)
-        # for name in embedding.choices:
-        #     updates = self._get_search_space_updates(prefix=name)
-        #     config_space = available_embedding[name].get_hyperparameter_search_space(
-        # dataset_properties,  # type: ignore
-        #                                                                              **updates)
-        #     parent_hyperparameter = {'parent': embedding, 'value': name}
-        #     cs.add_configuration_space(
-        #         name,
-        #         config_space,
-        #         parent_hyperparameter=parent_hyperparameter
-        #     )
-
-        # self.configuration_space_ = cs
-        # self.dataset_properties_ = dataset_properties
-        # return cs
 
     def transform(self, X: np.ndarray) -> np.ndarray:
         assert self.choice is not None, "Cannot call transform before the object is initialized"

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

@@ -334,6 +334,35 @@ def on_epoch_start(self, X: Dict[str, Any], epoch: int) -> None:
         """
         pass
 
+    def _swa_update(self) -> None:
+        """
+        perform swa model update
+        """
+        assert self.swa_model is not None, "SWA model can't be none when" \
+                                           " stochastic weight averaging is enabled"
+        self.swa_model.update_parameters(self.model)
+        self.swa_updated = True
+
+    def _se_update(self, epoch: int) -> None:
+        """
+        Add latest model or swa_model to model snapshot ensemble
+        Args:
+            epoch (int):
+                current epoch
+        """
+        assert self.model_snapshots is not None, "model snapshots container can't be " \
+                                                 "none when snapshot ensembling is enabled"
+        is_last_epoch = (epoch == self.budget_tracker.max_epochs)
+        if is_last_epoch and self.use_stochastic_weight_averaging:
+            model_copy = deepcopy(self.swa_model)
+        else:
+            model_copy = deepcopy(self.model)
+
+        assert model_copy is not None
+        model_copy.cpu()
+        self.model_snapshots.append(model_copy)
+        self.model_snapshots = self.model_snapshots[-self.se_lastk:]
+
     def on_epoch_end(self, X: Dict[str, Any], epoch: int) -> bool:
         """
         Optional place holder for AutoPytorch Extensions.
@@ -344,39 +373,19 @@ def on_epoch_end(self, X: Dict[str, Any], epoch: int) -> bool:
         if X['is_cyclic_scheduler']:
             if hasattr(self.scheduler, 'T_cur') and self.scheduler.T_cur == 0 and epoch != 1:
                 if self.use_stochastic_weight_averaging:
-                    assert self.swa_model is not None, "SWA model can't be none when" \
-                                                       " stochastic weight averaging is enabled"
-                    self.swa_model.update_parameters(self.model)
-                    self.swa_updated = True
+                    self._swa_update()
                 if self.use_snapshot_ensemble:
-                    assert self.model_snapshots is not None, "model snapshots container can't be " \
-                                                             "none when snapshot ensembling is enabled"
-                    is_last_epoch = (epoch == self.budget_tracker.max_epochs)
-                    if is_last_epoch and self.use_stochastic_weight_averaging:
-                        model_copy = deepcopy(self.swa_model)
-                    else:
-                        model_copy = deepcopy(self.model)
-
-                    assert model_copy is not None
-                    model_copy.cpu()
-                    self.model_snapshots.append(model_copy)
-                    self.model_snapshots = self.model_snapshots[-self.se_lastk:]
+                    self._se_update(epoch=epoch)
         else:
-            if epoch > self._budget_threshold:
-                if self.use_stochastic_weight_averaging:
-                    assert self.swa_model is not None, "SWA model can't be none when" \
-                                                       " stochastic weight averaging is enabled"
-                    self.swa_model.update_parameters(self.model)
-                    self.swa_updated = True
-                if self.use_snapshot_ensemble:
-                    assert self.model_snapshots is not None, "model snapshots container can't be " \
-                                                             "none when snapshot ensembling is enabled"
-                    model_copy = deepcopy(self.swa_model) if self.use_stochastic_weight_averaging \
-                        else deepcopy(self.model)
-                    assert model_copy is not None
-                    model_copy.cpu()
-                    self.model_snapshots.append(model_copy)
-                    self.model_snapshots = self.model_snapshots[-self.se_lastk:]
+            if epoch > self._budget_threshold and self.use_stochastic_weight_averaging:
+                self._swa_update()
+
+            if (
+                self.use_snapshot_ensemble
+                and self.budget_tracker.max_epochs is not None
+                and epoch > (self.budget_tracker.max_epochs - self.se_lastk)
+            ):
+                self._se_update(epoch=epoch)
         return False
 
     def _scheduler_step(