Update train.py

plot_confusion_matrix deprecated from scikit-learn 1.0 and removed in 1.2
edwenger · Jun 6, 2023 · 5a3d305 · 5a3d305
1 parent b07a0e6
commit 5a3d305
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -2,7 +2,7 @@
 import os
 
 from sklearn.ensemble import RandomForestClassifier
-from sklearn.metrics import plot_confusion_matrix
+from sklearn.metrics import ConfusionMatrixDisplay
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -29,5 +29,5 @@
     outfile.write(metrics)
 
 # Plot it
-disp = plot_confusion_matrix(clf, X_test, y_test, normalize="true", cmap=plt.cm.Blues)
+disp = ConfusionMatrixDisplay.from_estimator(clf, X_test, y_test, normalize="true", cmap=plt.cm.Blues)
 plt.savefig("plot.png")