Как построить тепловую карту на основе результатов GradientBoost?

Question

Я хочу напечатать тепловую карту матрицы путаницы, основываясь на моих результатах y_predict и y_train.

Я немного застрял и уже просмотрел документацию для панд тепловой карты, но все еще не знаю, как применить ее к своим результатам.Набор данных, который я использовал, касается доходов и содержит категориальные и числовые данные.Я уже применил классификатор GB и получил результат.Остается только тепловая карта.

print(confusion_matrix(y_train,y_pred_train))
print(y_train)

это был результат

Confusion Matrix:


[[14151   710]
 [ 1844  2831]]
Name: income, Length: 19536, dtype: int64

это была попытка создать тепловую карту

import seaborn as sns
class_names = y_train, y_pred_train

def print_confusion_matrix(confusion_matrix, class_names, figsize = (10,7), fontsize=14):
    df_cm = pd.DataFrame(
        confusion_matrix, index=class_names, columns=class_names, 
    )
    fig = plt.figure(figsize=figsize)
    try:
        heatmap = sns.heatmap(df_cm, annot=True, fmt="d")
    except ValueError:
        raise ValueError("Confusion matrix values must be integers.")
    heatmap.yaxis.set_ticklabels(heatmap.yaxis.get_ticklabels(), rotation=0, ha='right', fontsize=fontsize)
    heatmap.xaxis.set_ticklabels(heatmap.xaxis.get_ticklabels(), rotation=45, ha='right', fontsize=fontsize)
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    return fig

, которая вернулась

NameError                                 Traceback (most recent call last)
<ipython-input-36-3bd0e9ee90a4> in <module>()
     18     plt.xlabel('Predicted label')
     19     return fig
---> 20 print(fig)

NameError: name 'fig' is not defined

Так чего же мне не хватает, когда я делаю тепловую карту матрицы путаницы с моими результатами?

Answer 1

Вы можете вызвать функцию print_confusion_matrix со списком матриц смешения и именами классов в качестве параметров:

def print_confusion_matrix(confusion_matrix, class_names, figsize = (10,7), fontsize=14):
    df_cm = pd.DataFrame(
        confusion_matrix, index=class_names, columns=class_names, 
    )
    fig = plt.figure(figsize=figsize)
    try:
        heatmap = sns.heatmap(df_cm, annot=True, fmt="d")
    except ValueError:
        raise ValueError("Confusion matrix values must be integers.")
    heatmap.yaxis.set_ticklabels(heatmap.yaxis.get_ticklabels(), rotation=0, ha='right', fontsize=fontsize)
    heatmap.xaxis.set_ticklabels(heatmap.xaxis.get_ticklabels(), rotation=45, ha='right', fontsize=fontsize)
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    return fig

confusion_matrix = np.array([[14151, 710], [1844, 2831]])
fig = print_confusion_matrix(confusion_matrix, ['0', '1'])

Вывод: