Speedup similarity forest regressor and add parallelization at the extracted...

Speedup similarity forest regressor and add parallelization at the extracted forest size level in the training

code/bolsonaro/models/similarity_forest_regressor.py

-from bolsonaro.utils import tqdm_joblib
-
 from sklearn.ensemble import RandomForestRegressor
 from sklearn.metrics import mean_squared_error
 from sklearn.base import BaseEstimator
 from abc import abstractmethod, ABCMeta
 import numpy as np
-from joblib import Parallel, delayed
 from tqdm import tqdm
 
 
@@ -34,35 +31,40 @@ class SimilarityForestRegressor(BaseEstimator, metaclass=ABCMeta):
         selected_trees = list()
         tree_list = list(self._estimator.estimators_)
 
-        for _ in range(self._extracted_forest_size):
+        val_scores = list()
+        with tqdm(tree_list) as tree_pred_bar:
+            tree_pred_bar.set_description('[Initial tree predictions]')
+            for tree in tree_pred_bar:
+                val_scores.append(tree.predict(X_val))
+            tree_pred_bar.update(1)
+
+        with tqdm(range(self._extracted_forest_size), disable=False) as pruning_forest_bar:
+            pruning_forest_bar.set_description(f'[Pruning forest s={self._extracted_forest_size}]')
+            for i in pruning_forest_bar:
                 best_similarity = 100000
                 found_index = 0
-            for i in range(len(tree_list)):
-                lonely_tree = tree_list[i]
-                del tree_list[i]
-                begin_time = time.time()
-                with tqdm_joblib(tqdm(total=len(tree_list), disable=True)) as job_pb:
-                    val_list = Parallel(n_jobs=-1)(delayed(self._tree_predict_job)(
-                        job_pb, tree_list[i], X_val)
-                        for i in range(len(tree_list)))
-                val_list = np.array(val_list)
-                val_mean = np.mean(val_list, axis=0)
+                with tqdm(range(len(tree_list)), disable=False) as tree_list_bar:
+                    tree_list_bar.set_description(f'[Tree selection s={self._extracted_forest_size} #{i}]')
+                    for j in tree_list_bar:
+                        lonely_tree = tree_list[j]
+                        del tree_list[j]
+                        val_mean = np.mean(np.asarray(val_scores), axis=0)
                         val_score = self._score_metric(val_mean, y_val)
                         temp_similarity = abs(forest_pred - val_score)
                         if (temp_similarity < best_similarity):
-                    found_index = i
+                            found_index = j
                             best_similarity = temp_similarity
-                tree_list.insert(i, lonely_tree)
+                        tree_list.insert(j, lonely_tree)
+                        val_scores.insert(j, lonely_tree.predict(X_val))
+                        tree_list_bar.update(1)
                 selected_trees.append(tree_list[found_index])
                 del tree_list[found_index]
+                del val_scores[found_index]
+                pruning_forest_bar.update(1)
 
         pruned_forest = list(set(forest) - set(selected_trees))
         self._estimator.estimators_ = pruned_forest
 
-    def _tree_predict_job(self, job_pb, tree, X_val):
-        val_pred = tree.predict(X_val)
-        return val_pred
-
     def score(self, X, y):
         test_list = list()
         for mod in self._estimator.estimators_:

code/train.py

@@ -21,7 +21,7 @@ import numpy as np
 import shutil
 
 
-def process_job(seed, parameters, experiment_id, hyperparameters):
+def seed_job(seed_job_pb, seed, parameters, experiment_id, hyperparameters, verbose):
     """
     Experiment function.
 
@@ -34,7 +34,6 @@ def process_job(seed, parameters, experiment_id, hyperparameters):
     """
     logger = LoggerFactory.create(LOG_PATH, 'training_seed{}_ti{}'.format(
         seed, threading.get_ident()))
-    logger.info('seed={}'.format(seed))
 
     seed_str = str(seed)
     experiment_id_str = str(experiment_id)
@@ -55,9 +54,14 @@ def process_job(seed, parameters, experiment_id, hyperparameters):
     trainer = Trainer(dataset)
 
     if parameters['extraction_strategy'] != 'none':
-        for extracted_forest_size in parameters['extracted_forest_size']:
-            logger.info('extracted_forest_size={}'.format(extracted_forest_size))
-            sub_models_dir = models_dir + os.sep + 'extracted_forest_sizes' + os.sep + str(extracted_forest_size)
+        with tqdm_joblib(tqdm(total=len(parameters['extracted_forest_size']), disable=not verbose)) as extracted_forest_size_job_pb:
+            Parallel(n_jobs=-1)(delayed(extracted_forest_size_job)(extracted_forest_size_job_pb, parameters['extracted_forest_size'][i],
+                models_dir, seed, parameters, dataset, hyperparameters, experiment_id, trainer)
+                for i in range(len(parameters['extracted_forest_size'])))
+    else:
+        forest_size = hyperparameters['n_estimators']
+        logger.info('Base forest training with fixed forest size of {}'.format(forest_size))
+        sub_models_dir = models_dir + os.sep + 'forest_size' + os.sep + str(forest_size)
 
         # Check if the result file already exists
         already_exists = False
@@ -70,13 +74,11 @@ def process_job(seed, parameters, experiment_id, hyperparameters):
                     already_exists = os.path.getsize(os.path.join(sub_models_dir, file_name)) > 0
                     break
         if already_exists:
-                logger.info(f'Extracted forest {extracted_forest_size} result already exists. Skipping...')
-                continue
-
+            logger.info('Base forest result already exists. Skipping...')
+        else:
             pathlib.Path(sub_models_dir).mkdir(parents=True, exist_ok=True)
-
             model_parameters = ModelParameters(
-                extracted_forest_size=extracted_forest_size,
+                extracted_forest_size=forest_size,
                 normalize_D=parameters['normalize_D'],
                 subsets_used=parameters['subsets_used'],
                 normalize_weights=parameters['normalize_weights'],
@@ -91,10 +93,17 @@ def process_job(seed, parameters, experiment_id, hyperparameters):
             trainer.init(model, subsets_used=parameters['subsets_used'])
             trainer.train(model)
             trainer.compute_results(model, sub_models_dir)
-    else:
-        forest_size = hyperparameters['n_estimators']
-        logger.info('Base forest training with fixed forest size of {}'.format(forest_size))
-        sub_models_dir = models_dir + os.sep + 'forest_size' + os.sep + str(forest_size)
+    logger.info(f'Training done for seed {seed_str}')
+    seed_job_pb.update(1)
+
+def extracted_forest_size_job(extracted_forest_size_job_pb, extracted_forest_size, models_dir,
+    seed, parameters, dataset, hyperparameters, experiment_id, trainer):
+
+    logger = LoggerFactory.create(LOG_PATH, 'training_seed{}_extracted_forest_size{}_ti{}'.format(
+        seed, extracted_forest_size, threading.get_ident()))
+    logger.info('extracted_forest_size={}'.format(extracted_forest_size))
+
+    sub_models_dir = models_dir + os.sep + 'extracted_forest_sizes' + os.sep + str(extracted_forest_size)
 
     # Check if the result file already exists
     already_exists = False
@@ -102,17 +111,18 @@ def process_job(seed, parameters, experiment_id, hyperparameters):
         sub_models_dir_files = os.listdir(sub_models_dir)
         for file_name in sub_models_dir_files:
             if '.pickle' != os.path.splitext(file_name)[1]:
-                    continue
+                return
             else:
                 already_exists = os.path.getsize(os.path.join(sub_models_dir, file_name)) > 0
                 break
     if already_exists:
-            logger.info('Base forest result already exists. Skipping...')
-        else:
-            pass
+        logger.info(f'Extracted forest {extracted_forest_size} result already exists. Skipping...')
+        return
+
     pathlib.Path(sub_models_dir).mkdir(parents=True, exist_ok=True)
+
     model_parameters = ModelParameters(
-                extracted_forest_size=forest_size,
+        extracted_forest_size=extracted_forest_size,
         normalize_D=parameters['normalize_D'],
         subsets_used=parameters['subsets_used'],
         normalize_weights=parameters['normalize_weights'],
@@ -127,7 +137,6 @@ def process_job(seed, parameters, experiment_id, hyperparameters):
     trainer.init(model, subsets_used=parameters['subsets_used'])
     trainer.train(model)
     trainer.compute_results(model, sub_models_dir)
-    logger.info('Training done')
 
 """
 Command lines example for stage 1:
@@ -287,6 +296,6 @@ if __name__ == "__main__":
             )
 
     # Run as much job as there are seeds
-    with tqdm_joblib(tqdm(total=len(seeds), disable=not args.verbose)) as progress_bar:
-        Parallel(n_jobs=args.job_number)(delayed(process_job)(seeds[i],
-            parameters, experiment_id, hyperparameters) for i in range(len(seeds)))
+    with tqdm_joblib(tqdm(total=len(seeds), disable=not args.verbose)) as seed_job_pb:
+        Parallel(n_jobs=args.job_number)(delayed(seed_job)(seed_job_pb, seeds[i],
+            parameters, experiment_id, hyperparameters, args.verbose) for i in range(len(seeds)))