Improve confusion matrix plot (#5273)

* update Signed-off-by: Weichen Xu <weichen.xu@databricks.com> * fix Signed-off-by: Weichen Xu <weichen.xu@databricks.com> * update Signed-off-by: Weichen Xu <weichen.xu@databricks.com>
mlflow · Jan 17, 2022 · 847eb6b · 847eb6b
1 parent 00921fc
commit 847eb6b
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
diff --git a/mlflow/models/evaluation/default_evaluator.py b/mlflow/models/evaluation/default_evaluator.py
@@ -599,10 +599,18 @@ def _evaluate_classifier(self):
         )
 
         def plot_confusion_matrix():
-            sk_metrics.ConfusionMatrixDisplay(
-                confusion_matrix=confusion_matrix,
-                display_labels=self.label_list,
-            ).plot(cmap="Blues")
+            import matplotlib
+
+            with matplotlib.rc_context(
+                {
+                    "font.size": min(10, 50.0 / self.num_classes),
+                    "axes.labelsize": 10,
+                }
+            ):
+                sk_metrics.ConfusionMatrixDisplay(
+                    confusion_matrix=confusion_matrix,
+                    display_labels=self.label_list,
+                ).plot(cmap="Blues")
 
         if hasattr(sk_metrics, "ConfusionMatrixDisplay"):
             self._log_image_artifact(

diff --git a/mlflow/sklearn/utils.py b/mlflow/sklearn/utils.py
@@ -288,10 +288,22 @@ def _get_classifier_artifacts(fitted_estimator, prefix, X, y_true, sample_weight
     if not _is_plotting_supported():
         return []
 
+    def plot_confusion_matrix(*args, **kwargs):
+        import matplotlib
+
+        num_classes = len(set(y_true))
+        with matplotlib.rc_context(
+            {
+                "font.size": min(10.0, 50.0 / num_classes),
+                "axes.labelsize": 10.0,
+            }
+        ):
+            return sklearn.metrics.plot_confusion_matrix(*args, **kwargs)
+
     classifier_artifacts = [
         _SklearnArtifact(
             name=prefix + "confusion_matrix",
-            function=sklearn.metrics.plot_confusion_matrix,
+            function=plot_confusion_matrix,
             arguments=dict(
                 estimator=fitted_estimator,
                 X=X,