diff --git a/code/bolsonaro/data/dataset.py b/code/bolsonaro/data/dataset.py index 7108eb5781e1ef2b69926e9fd2239deaa81f44e2..3bca3ce2477a3db8934159bd786eebb78c8d236e 100644 --- a/code/bolsonaro/data/dataset.py +++ b/code/bolsonaro/data/dataset.py @@ -14,10 +14,6 @@ class Dataset(object): def task(self): return self._task - @property - def dataset_parameters(self): - return self._dataset_parameters - @property def X_train(self): return self._X_train diff --git a/code/bolsonaro/trainer.py b/code/bolsonaro/trainer.py index e1bc893dca0dae03c3b24a3265868547004b3a3e..10eea289a2989765fd66bfea5ae7fa940d69de3e 100644 --- a/code/bolsonaro/trainer.py +++ b/code/bolsonaro/trainer.py @@ -2,6 +2,7 @@ from bolsonaro.models.model_raw_results import ModelRawResults from bolsonaro.models.omp_forest_regressor import OmpForestRegressor from bolsonaro.models.omp_forest_classifier import OmpForestBinaryClassifier, OmpForestMulticlassClassifier from bolsonaro.error_handling.logger_factory import LoggerFactory +from bolsonaro.data.task import Task from . import LOG_PATH from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier @@ -16,13 +17,30 @@ class Trainer(object): Class capable of fitting any model object to some prepared data then evaluate and save results through the `train` method. """ - def __init__(self, dataset): + def __init__(self, dataset, regression_score_metric=mean_squared_error, classification_score_metric=accuracy_score, + base_regression_score_metric=mean_squared_error, base_classification_score_metric=accuracy_score): """ :param dataset: Object with X_train, y_train, X_dev, y_dev, X_test and Y_test attributes """ self._dataset = dataset self._logger = LoggerFactory.create(LOG_PATH, __name__) + self._regression_score_metric = regression_score_metric + self._classification_score_metric = classification_score_metric + self._base_regression_score_metric = base_regression_score_metric + self._base_classification_score_metric = base_classification_score_metric + self._score_metric_name = regression_score_metric.__name__ if dataset.task == Task.REGRESSION \ + else classification_score_metric.__name__ + self._base_score_metric_name = base_regression_score_metric.__name__ if dataset.task == Task.REGRESSION \ + else base_classification_score_metric.__name__ + + @property + def score_metric_name(self): + return self._score_metric_name + + @property + def base_score_metric_name(self): + return self._base_score_metric_name def init(self, model): if type(model) in [RandomForestRegressor, RandomForestClassifier]: @@ -75,27 +93,25 @@ class Trainer(object): def __score_func(self, model, X, y_true): if type(model) in [OmpForestRegressor, RandomForestRegressor]: y_pred = model.predict(X) - result = mean_squared_error(y_true, y_pred) + result = self._regression_score_metric(y_true, y_pred) elif type(model) in [OmpForestBinaryClassifier, OmpForestMulticlassClassifier, RandomForestClassifier]: y_pred = model.predict(X) - result = accuracy_score(y_true, y_pred) - + result = self._classification_score_metric(y_true, y_pred) return result def __score_func_base(self, model, X, y_true): if type(model) == OmpForestRegressor: y_pred = model.predict_base_estimator(X) - result = mean_squared_error(y_true, y_pred) + result = self._base_regression_score_metric(y_true, y_pred) elif type(model) in [OmpForestBinaryClassifier, OmpForestMulticlassClassifier]: y_pred = model.predict_base_estimator(X) - result = accuracy_score(y_true, y_pred) + result = self._base_classification_score_metric(y_true, y_pred) elif type(model) == RandomForestClassifier: y_pred = model.predict(X) - result = accuracy_score(y_true, y_pred) + result = self._base_classification_score_metric(y_true, y_pred) elif type(model) == RandomForestRegressor: y_pred = model.predict(X) - result = mean_squared_error(y_true, y_pred) - + result = self._base_regression_score_metric(y_true, y_pred) return result def compute_results(self, model, models_dir): @@ -113,8 +129,8 @@ class Trainer(object): train_score_base=self.__score_func_base(model, self._dataset.X_train, self._dataset.y_train), dev_score_base=self.__score_func_base(model, self._dataset.X_dev, self._dataset.y_dev), test_score_base=self.__score_func_base(model, self._dataset.X_test, self._dataset.y_test), - score_metric='mse' if type(model) in [RandomForestRegressor, RandomForestClassifier] \ - else model.DEFAULT_SCORE_METRIC, # TODO: resolve the used metric in a proper way + score_metric=self._score_metric_name, + base_score_metric=self._base_score_metric_name ) results.save(models_dir) self._logger.info("Base performance on test: {}".format(results.test_score_base))