tensorflow2 训练mnist 显示卷积核及卷积层的feature map

保存一下代码为 test.py, 在同级目录下新建文件夹 models和 logs。注意修改加载数据集的路径，若不写 path参数，则从网上下载，比较慢，建议先将数据集从百度网盘上下载下来。第一次要先训练，模型保存在 models文件夹下。

python test.py --train

有已经训练好的模型，就不用训练了

python test.py 

代码：

from __future__ import absolute_import, division, print_function, unicode_literals

# 安装 TensorFlow

import tensorflow as tf
import tensorflow_datasets as tfds 
import matplotlib.pyplot as plt
import numpy as np
import time
from datetime import datetime
import math
import io
import sys, getopt

from tensorflow.keras import backend as K
from tensorflow.keras.models import load_model

#返回数据维度(samples, width, height, channel): (60000, 28, 28, 1)
def loadData():
    mnist = tf.keras.datasets.mnist

    # 一定要是绝对路径， 若不加 path参数，就从网上下载
    # 加载后的数据维数是 (60000, 28, 28)。 60000 个 28 x 28的图片
    (x_train, y_train), (x_test, y_test) = mnist.load_data(path='E:\\AI\\tensorflow2study\\data\\mnist.npz')
    x_train, x_test = x_train / 255.0, x_test / 255.0

    # 增加channel ：(60000, 28, 28) -> (60000, 28, 28, 1)，为了符合模型输入要求
    x_train = x_train[:, :, :, np.newaxis]
    x_test = x_test[:, :, :, np.newaxis]

    return (x_train, y_train), (x_test, y_test)

# 建模
def buildModel():
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Conv2D(6, (5, 5), activation='relu', input_shape=(28, 28, 1)))
    model.add(tf.keras.layers.MaxPooling2D((2, 2)))
    model.add(tf.keras.layers.Conv2D(16, (3, 3), activation='relu'))
    model.add(tf.keras.layers.MaxPooling2D((2, 2)))
    model.add(tf.keras.layers.Flatten(input_shape=(28, 28)))
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(128, activation='relu'))
    model.add(tf.keras.layers.Dense(10, activation='softmax'))

    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    # 打印模型结构
    model.summary()
        
    return model


# 输入 images: [num, width, height]
def image_grid(images, labels):

    # Create a figure to contain the plot.
    images_num = len(images)
    col_num = 5
    row_num = math.ceil(images_num / col_num)

    fig_width = col_num * 2
    fig_height = row_num * 2
    figure = plt.figure(figsize=(fig_width, fig_height))
    
    for i in range(images_num):
        # Start next subplot.
        plt.subplot(row_num, col_num, i + 1, title = labels[i])#每行5个
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        plt.imshow(images[i], cmap = plt.cm.binary)

    # plt.ion();
    # plt.show();
    # plt.pause(0)
    figure.tight_layout(pad=0.4, w_pad=3.0, h_pad=3.0)
    return figure

# 将figure中的图像转换成 tf 图像
def plot_to_image(figure):
    """Converts the matplotlib plot specified by 'figure' to a PNG image and
    returns it. The supplied figure is closed and inaccessible after this call."""
    # Save the plot to a PNG in memory.
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    # Closing the figure prevents it from being displayed directly inside
    # the notebook.
    plt.close(figure)
    buf.seek(0)
    # Convert PNG buffer to TF image
    image = tf.image.decode_png(buf.getvalue(), channels=4)
    # Add the batch dimension
    image = tf.expand_dims(image, 0)
    return image

def saveTrainInputImageToTf(train_images, train_labels):
    # Prepare the plot
    figure = image_grid(train_images, train_labels)
    # plot_to_image(figure)
    logdir = "logs/plots/" + datetime.now().strftime("%Y%m%d-%H%M%S")
    file_writer = tf.summary.create_file_writer(logdir)
    # # Convert to image and log
    with file_writer.as_default():
        tf.summary.image("Training data", plot_to_image(figure), step=0)


def saveTrainKernalToTf(model):
    weights = model.weights
    for i in range(len(model.weights)):
        # 只有是卷积核才显示，卷积核的维度是4 (kernal_size, kernal_size, inputs_len, outputs_len)
        if len(model.weights[i].shape) != 4:
            continue

        inputs_len = model.weights[i].shape[2]
        outputs_len = model.weights[i].shape[3]

        # 提取 inputs_len * outputs_len 个卷积核
        kernel_array = []
        kernal_label_array = []
        for j in range(inputs_len):
            for k in range(outputs_len):
                kernel = weights[i][:, :, j, k]
                kernel_array.append(kernel)
                kernal_label_array.append(str(j) + "_" + str(k)+ "_kernal")

        kernel_array = np.array(kernel_array)

        print("Kernal Shape: ", kernel_array.shape)

        figure = image_grid(kernel_array, kernal_label_array)
        # plot_to_image(figure)
        logdir = "logs/plots/" + datetime.now().strftime("%Y%m%d-%H%M%S")
        file_writer = tf.summary.create_file_writer(logdir)
        # # Convert to image and log
        with file_writer.as_default():
            tf.summary.image("Training kernal", plot_to_image(figure), step=0)


def saveConvImage(model, train_images, train_labels):
    for layer in model.layers:
        if 'conv' in layer.name:
            print(layer.name)

            activations_f = K.function([model.layers[0].input], [layer.output])
            activations = activations_f((train_images[0:2, :, :, :], False))#选择前面两个样本
            
            activations = np.array(activations)
            print(activations.shape) # (1, 2, 24, 24, 6) （none, 样本个数，特征图尺寸长，特征图尺寸宽，特征图个数）

            conv_layer = []
            conv_layer_label = []
            for sample_index in range(activations.shape[1]):
                for i in range(activations.shape[4]):# 特征图个数
                    # activations[0, sample_index, :, :, i] 为第 sample_index 个样本的 第i个特征图
                    conv_layer.append(activations[0, sample_index, :, :, i])
                    conv_layer_label.append(train_labels[sample_index])

            conv_layer = np.array(conv_layer)
            print(conv_layer.shape)
            figure = image_grid(conv_layer, conv_layer_label)
            # plot_to_image(figure)
            logdir = "logs/plots/" + datetime.now().strftime("%Y%m%d-%H%M%S")
            file_writer = tf.summary.create_file_writer(logdir)
            # # Convert to image and log
            with file_writer.as_default():
                tf.summary.image("conv layer", plot_to_image(figure), step=0)



if __name__ == "__main__":
    (train_images, train_labels), (test_images, test_labels) = loadData()

    # 绘制输入图
    saveTrainInputImageToTf(train_images[0:25, :, :, 0], train_labels[0:25])

    # 命令行获取参数，是训练模型还是加载已训练好的模型
    is_train_model = False
    try:
        opts, args = getopt.getopt(sys.argv[1:], "h", ["train"])
    except getopt.GetoptError:
        print('test.py --train')
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print('test.py --train')
            sys.exit()
        elif opt in ("--train"):
            is_train_model = True

    if is_train_model:
        model = buildModel()
        for i in range(5):
            # Train the classifier.
            model.fit(
                train_images,
                train_labels,
                epochs=5,
                verbose=0, # Suppress chatty output
                callbacks=[
                    tf.keras.callbacks.TensorBoard(
                        log_dir=".\\logs"
                    )
                ],
                validation_data=(test_images, test_labels),
            )
        model.save(".\\models\\my_model.h5")
    else:
        model = load_model(".\\models\\my_model.h5")

    # 绘制训练完成后的卷积核
    saveTrainKernalToTf(model)
    # 显示经过卷积层的 feature map
    saveConvImage(model, train_images, train_labels)

程序结束后，开启 tensorboard

tensorboard --logdir logs

打开浏览器，输入http://localhost:6006, 点击 images tab:

前 25 个样本显示：

第一个卷积层（1x6个大小为5x5 的卷积核），第二个卷积层（6x16个大小为3x3 的卷积核）显示：

两个样本（5和0）两个卷积层的 feature map显示：