dont plot training scores for AUC in CatBoost (mljar#360)

ajschumacher · Mar 31, 2021 · c461107 · c461107
1 parent ff82af8
commit c461107
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 19 deletions.
diff --git a/supervised/algorithms/catboost.py b/supervised/algorithms/catboost.py
@@ -230,23 +230,30 @@ def fit(
             self.best_ntree_limit = self.model.tree_count_
 
         if log_to_file is not None:
-            train_scores = self.model.evals_result_["learn"][self.log_metric_name]
-            validation_scores = self.model.evals_result_["validation"][
+            train_scores = self.model.evals_result_["learn"].get(self.log_metric_name)
+            validation_scores = self.model.evals_result_["validation"].get(
                 self.log_metric_name
-            ]
+            )
             if model_init is not None:
-                train_scores = (
-                    model_init.evals_result_["learn"][self.log_metric_name]
-                    + train_scores
-                )
-                validation_scores = (
-                    model_init.evals_result_["validation"][self.log_metric_name]
-                    + validation_scores
-                )
+                if train_scores is not None:
+                    train_scores = (
+                        model_init.evals_result_["learn"].get(self.log_metric_name)
+                        + train_scores
+                    )
+                if validation_scores is not None:
+                    validation_scores = (
+                        model_init.evals_result_["validation"].get(self.log_metric_name)
+                        + validation_scores
+                    )
+            iteration = None 
+            if train_scores is not None:
+                iteration = range(len(validation_scores))
+            elif validation_scores is not None:
+                iteration = range(len(validation_scores))
 
             result = pd.DataFrame(
                 {
-                    "iteration": range(len(train_scores)),
+                    "iteration": iteration,
                     "train": train_scores,
                     "validation": validation_scores,
                 }
@@ -295,6 +302,8 @@ def get_metric_name(self):
             return None
         if metric == "Logloss":
             return "logloss"
+        elif metric == "AUC":
+            return "auc"
         elif metric == "MultiClass":
             return "logloss"
         elif metric == "RMSE":

diff --git a/supervised/utils/learning_curves.py b/supervised/utils/learning_curves.py
@@ -88,13 +88,15 @@ def plot_iterations(
             repeat_str = f" Reapeat {repeat+1}," if repeat is not None else ""
             # if trees_in_iteration is not None:
             #    df.iteration = df.iteration * trees_in_iteration
-            plt.plot(
-                df.iteration,
-                df.train,
-                "--",
-                color=colors[fold],
-                label=f"Fold {fold+1},{repeat_str} train",
-            )
+            any_none = np.sum(pd.isnull(df.train))
+            if any_none == 0:
+                plt.plot(
+                    df.iteration,
+                    df.train,
+                    "--",
+                    color=colors[fold],
+                    label=f"Fold {fold+1},{repeat_str} train",
+                )
             any_none = np.sum(pd.isnull(df.test))
             if any_none == 0:
                 plt.plot(