带有不同标签的轴的混淆矩阵

时间:2019-03-18 18:58:12

标签: python scikit-learn

我试图训练和测试Naive Bayes分类器。

以下是我的代码的一部分:

from sklearn.feature_extraction.text import CountVectorizer
matrix = CountVectorizer(ngram_range=(1,1))
X = matrix.fit_transform(data).toarray()
y =  [re.sub('[^A-Za-z]', ' ', y).strip(' ') for y in mobiles.iloc[:, 2]]

# split train and test data
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y)

# Naive Bayes 
from sklearn.naive_bayes import GaussianNB

classifier = GaussianNB()

classifier.fit(X_train, y_train)

# predict class
y_pred = classifier.predict(X_test)

res = pd.DataFrame({'y_test':y_test, 'y_pred':y_pred})
print(res)
# Confusion matrix
from sklearn.metrics import confusion_matrix, classification_report, accuracy_score, f1_score
cm = confusion_matrix(y_test, y_pred)
cr = classification_report(y_test, y_pred)

brands = list(set(y))

accuracy = accuracy_score(y_test, y_pred)
print("accuracy:", accuracy)
print("Confusion Matrix:")
import seaborn as sn
import pandas  as pd
import matplotlib.pyplot as plt

aylabels = brands #[str(i) for i in aylabels]
axlabels = brands #[str(i) for i in range(50)]

plt.figure(figsize=(10, 10))
sn.set(font_scale=1.4)  # for label size
sn.heatmap(cm, annot=True, annot_kws={"size": 12}, xticklabels=axlabels, yticklabels=aylabels)  # font size
plt.show()

以下是我在上面的代码中根据hitmap构建的混淆矩阵。虽然我将aylabelsaxlabels设置为相同,但是图中的行和列不同。

enter image description here

我不知道发生了什么事!

1 个答案:

答案 0 :(得分:1)

您忘记在confusion_matrix中设置标签参数。

cm = confusion_matrix(y_test, y_pred, labels=brands)