Merge pull request #49 from m3dev/add_lgbm_mode_factorys

Add lgbm and catboost model factories

Author: vaaaaanquish

redshells/factory/optuna_param_factory.py

@@ -28,10 +28,45 @@ def _xgbclassifier_default(trial: optuna.trial.Trial):
     return param
 
 
+def _lgbmclassifier_default(trial: optuna.trial.Trial):
+    # TODO: using LightGBMTuner
+    params = {
+        'boosting_type': trial.suggest_categorical('boosting', ['gbdt', 'dart', 'goss']),
+        'objective': 'binary',
+        'metric': ['binary', 'binary_error', 'auc'],
+        'num_leaves': trial.suggest_int("num_leaves", 10, 500),
+        'learning_rate': trial.suggest_loguniform("learning_rate", 1e-5, 1),
+        'feature_fraction': trial.suggest_uniform("feature_fraction", 0.0, 1.0),
+    }
+    if params['boosting_type'] == 'dart':
+        params['drop_rate'] = trial.suggest_loguniform('drop_rate', 1e-8, 1.0)
+        params['skip_drop'] = trial.suggest_loguniform('skip_drop', 1e-8, 1.0)
+    if params['boosting_type'] == 'goss':
+        params['top_rate'] = trial.suggest_uniform('top_rate', 0.0, 1.0)
+        params['other_rate'] = trial.suggest_uniform('other_rate', 0.0, 1.0 - params['top_rate'])
+
+    return params
+
+
+def _catboostclassifier_default(trial: optuna.trial.Trial):
+    params = {
+        'iterations': trial.suggest_int('iterations', 50, 300),
+        'depth': trial.suggest_int('depth', 4, 10),
+        'learning_rate': trial.suggest_loguniform('learning_rate', 0.01, 0.3),
+        'random_strength': trial.suggest_int('random_strength', 0, 100),
+        'bagging_temperature': trial.suggest_loguniform('bagging_temperature', 0.01, 100.00),
+        'od_type': trial.suggest_categorical('od_type', ['IncToDec', 'Iter']),
+        'od_wait': trial.suggest_int('od_wait', 10, 50)
+    }
+
+    return params
+
 class _OptunaParamFactory(metaclass=Singleton):
     def __init__(self):
         self._rules = dict()
         self._rules['XGBClassifier_default'] = _xgbclassifier_default
+        self._rules['LGBMClassifier_default'] = _lgbmclassifier_default
+        self._rules['CatBoostClassifier_default'] = _catboostclassifier_default
 
     def get(self, key: str, trial: optuna.trial.Trial):
         if key not in self._rules:

redshells/factory/prediction_model_factory.py

@@ -18,6 +18,18 @@ def __init__(self):
         except ImportError:
             pass
 
+        try:
+            import lightgbm
+            self._models['LGBMClassifier'] = lightgbm.LGBMClassifier
+        except ImportError:
+            pass
+
+        try:
+            import catboost
+            self._models['CatBoostClassifier'] = catboost.CatBoostClassifier
+        except ImportError:
+            pass
+
     def get(self, key: str):
         if key in self._models:
             return self._models[key]