python中使用keras训练mnist模型，用yellowbrick绘制模型的confusion_matrix（混淆矩阵）可视化

代码如下：

import lmdb
import numpy as np
import tensorflow as tf
import binascii
import caffe
import matplotlib.pyplot as plt
import PIL.Image
from io import BytesIO

import keras
from keras import Sequential
from keras.layers import Dense, Flatten, Dropout, MaxPool2D, Conv2D
from keras.layers.normalization import BatchNormalization

from sklearn.metrics import confusion_matrix
from sklearn.preprocessing import LabelEncoder
from yellowbrick.style.palettes import color_sequence
from yellowbrick.style import find_text_color


cmap = color_sequence("YlOrRd")

with open('/workspace/jobs/20200114-013226-34ca/val.txt', 'r') as f:
    l = f.readlines()    
    
y_test = []
for i in l:
    y_test.append(int(i.split(' ')[1][0]))
    
y = np.array(y_test)
labels = np.unique(y)

env_db = lmdb.open('/workspace/jobs/20200114-013226-34ca/val_db')

input_ids = []
input_data = []

with env_db.begin() as txn:
    cursor = txn.cursor()
    
    for key, value in cursor:
        datum = caffe.proto.caffe_pb2.Datum()
        datum.ParseFromString(value)
        s = BytesIO()
        s.write(datum.data)
        s.seek(0)
        img = PIL.Image.open(s)
        img = np.array(img)
        img = img.flatten()
        
        input_ids.append(key)
        input_data.append(img)

y_train = np.asarray(input_data)
y_train = y_train.reshape(-1, 28, 28, 1)

# keras mnist network
model = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3),activation='relu',kernel_initializer='he_normal',input_shape=(28,28,1)))
model.add(Conv2D(32, kernel_size=(3, 3),activation='relu',kernel_initializer='he_normal'))
model.add(MaxPool2D((2, 2)))
model.add(Dropout(0.20))
model.add(Conv2D(64, (3, 3), activation='relu',padding='same',kernel_initializer='he_normal'))
model.add(Conv2D(64, (3, 3), activation='relu',padding='same',kernel_initializer='he_normal'))
model.add(MaxPool2D(pool_size=(2, 2)))
model.add(Dropout(0.25))
model.add(Conv2D(128, (3, 3), activation='relu',padding='same',kernel_initializer='he_normal'))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(BatchNormalization())
model.add(Dropout(0.25))

model.add(Dense(10, activation='softmax'))

model.compile(loss='mse',
              optimizer=keras.optimizers.RMSprop(),
              metrics=['accuracy'])

y2 = keras.utils.to_categorical(y, 10)
model.fit(y_train, y2, epochs=100, batch_size=32)

y_pred = model.predict_classes(y_train)

yp = np.asarray(y)
if yp.dtype.kind in {"i", "u"}:
    idx = yp
else:
    idx = LabelEncoder().fit_transform(yp)
y_true = np.asarray(labels)[idx]

yp = np.asarray(y_pred)
if yp.dtype.kind in {"i", "u"}:
    idx = yp
else:
    idx = LabelEncoder().fit_transform(yp)
y_pred = np.asarray(labels)[idx]

c_m = confusion_matrix(y_true, y_pred, labels=labels)

print("c_m >>>>>>> \n", c_m)

class_count = dict(zip(*np.unique(y_true, return_counts=True)))
selected_class_counts = []
for c in labels:
    try:
        selected_class_counts.append(class_count[c])
    except KeyError:
        selected_class_counts.append(0)
class_counts = np.array(selected_class_counts)

cm_display = c_m[::-1, ::]
n_classes = len(labels)
X, Y = np.arange(n_classes + 1), np.arange(n_classes + 1)

ax = plt.gca()

ax.set_ylim(bottom=0, top=cm_display.shape[0])
ax.set_xlim(left=0, right=cm_display.shape[1])

xticklabels = labels
yticklabels = labels[::-1]
ticks = np.arange(n_classes) + 0.5

ax.set(xticks=ticks, yticks=ticks)
ax.set_xticklabels(xticklabels, rotation="vertical")
ax.set_yticklabels(yticklabels)

edgecolors = []

for x in X[:-1]:
    for y in Y[:-1]:
        value = cm_display[x, y]
        svalue = "{:0.0f}".format(value)

        base_color = cmap(value / cm_display.max())
        text_color = find_text_color(base_color)

        if cm_display[x, y] == 0:
            text_color = "0.75"

        cx, cy = x + 0.5, y + 0.5
        ax.text(
            cy,
            cx,
            svalue,
            va="center",
            ha="center",
            color=text_color
        )
        lc = "k" if xticklabels[x] == yticklabels[y] else "w"
        edgecolors.append(lc)

vmin = 0.00001
vmax = cm_display.max()

ax.pcolormesh(
    X,
    Y,
    cm_display,
    vmin=vmin,
    vmax=vmax,
    edgecolor=edgecolors,
    cmap=cmap,
    linewidth="0.01",
)

# plt.show()
plt.savefig("/data/domon/confusion_matrix.png")

其实yellowbrick中有api可以直接实现，以上代码为拆解源码，实现自定义网络与自定义数据集来实现混淆矩阵的模型可视化。

yellowbrick实现如下：

from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split as tts
from sklearn.linear_model import LogisticRegression
from yellowbrick.classifier import ConfusionMatrix

# 将特征为8*8的像素图像转为单个特征数组X，标签为y
digits = load_digits()
X = digits.data
y = digits.target

X_train, X_test, y_train, y_test = tts(X, y, test_size =0.2, random_state=11)

model = LogisticRegression(multi_class="auto", solver="liblinear")

# 创建混淆矩阵模型
cm = ConfusionMatrix(model, classes=[0,1,2,3,4,5,6,7,8,9])

# 训练模型
cm.fit(X_train, y_train)

# score方法内对测试数据进行预测，并从scikit-learn创建confusion_matrix图像
cm.score(X_test, y_test)

# 显示图片
cm.show()

yellowbrick混淆矩阵可视化结果：

自定义-混淆矩阵可视化结果（模型 epochs=50，测试图片数量：250张）：

有需要改进的地方请提示，谢谢！