ADD allow changing ensemble size in fit_ensemble

zhangweijiqn · Apr 25, 2016 · b76a2b8 · b76a2b8
1 parent 9fd26ea
commit b76a2b8
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 11 deletions.
diff --git a/autosklearn/automl.py b/autosklearn/automl.py
@@ -468,18 +468,21 @@ def predict_proba(self, X):
         predictions = self.ensemble_.predict(all_predictions)
         return predictions
 
-    def fit_ensemble(self, task, metric, precision, dataset_name):
+    def fit_ensemble(self, task=None, metric=None, precision='32',
+                     dataset_name=None, ensemble_nbest=None,
+                     ensemble_size=None):
         if self._logger is None:
             self._logger = self._get_logger(dataset_name)
-        self._proc_ensemble = self._get_ensemble_process(1, task, metric,
-                                                         precision,
-                                                         dataset_name,
-                                                         max_iterations=1)
+
+        self._proc_ensemble = self._get_ensemble_process(
+            1, task, metric, precision, dataset_name, max_iterations=1,
+            ensemble_nbest=ensemble_nbest, ensemble_size=ensemble_size)
         self._proc_ensemble.main()
 
     def _get_ensemble_process(self, time_left_for_ensembles,
                               task=None, metric=None, precision=None,
-                              dataset_name=None, max_iterations=-1):
+                              dataset_name=None, max_iterations=-1,
+                              ensemble_nbest=None, ensemble_size=None):
 
         if task is None:
             task = self._task
@@ -489,15 +492,19 @@ def _get_ensemble_process(self, time_left_for_ensembles,
             precision = self.precision
         if dataset_name is None:
             dataset_name = self._dataset_name
+        if ensemble_nbest is None:
+            ensemble_nbest = self._ensemble_nbest
+        if ensemble_size is None:
+            ensemble_size = self._ensemble_size
 
         return EnsembleBuilder(autosklearn_tmp_dir=self._tmp_dir,
                                dataset_name=dataset_name,
                                task_type=task,
                                metric=metric,
                                limit=time_left_for_ensembles,
                                output_dir=self._output_dir,
-                               ensemble_size=self._ensemble_size,
-                               ensemble_nbest=self._ensemble_nbest,
+                               ensemble_size=ensemble_size,
+                               ensemble_nbest=ensemble_nbest,
                                seed=self._seed,
                                shared_mode=self._shared_mode,
                                precision=precision,

diff --git a/example/example_parallel.py b/example/example_parallel.py
@@ -26,7 +26,7 @@ def spawn_classifier(seed, dataset_name):
     X_test = X[1000:]
     y_test = y[1000:]
 
-    automl = AutoSklearnClassifier(time_left_for_this_task=180,
+    automl = AutoSklearnClassifier(time_left_for_this_task=60,
                                    per_run_time_limit=60,
                                    ml_memory_limit=1024,
                                    shared_mode=True,
@@ -64,16 +64,20 @@ def spawn_classifier(seed, dataset_name):
                                    per_run_time_limit=15,
                                    ml_memory_limit=1024,
                                    shared_mode=True,
+                                   ensemble_size=50,
+                                   ensemble_nbest=200,
                                    tmp_folder=tmp_folder,
                                    output_folder=output_folder,
-                                   ensemble_size=50,
                                    initial_configurations_via_metalearning=0,
                                    seed=1)
 
+    # Both the ensemble_size and ensemble_nbest parameters can be changed later
     automl.fit_ensemble(task=MULTICLASS_CLASSIFICATION,
                         metric=ACC_METRIC,
                         precision='32',
-                        dataset_name='digits')
+                        dataset_name='digits',
+                        ensemble_size=10,
+                        ensemble_nbest=10)
 
     predictions = automl.predict(X_test)
     print(automl.show_models())