mnist多个数字显示在一张图片并保存图片

import os
import scipy
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

data_path = r"D:\PycharmProjects\dataset"

def load_mnist(is_training=True):
    path = os.path.join(data_path, 'mnist')
    if is_training:
        fd = open(os.path.join(path, 'train-images-idx3-ubyte'))
        loaded = np.fromfile(file=fd, dtype=np.uint8)
        trainX = loaded[16:].reshape((60000, 28, 28, 1)).astype(np.float32)

        fd = open(os.path.join(path, 'train-labels-idx1-ubyte'))
        loaded = np.fromfile(file=fd, dtype=np.uint8)
        trY = loaded[8:].reshape((60000)).astype(np.int32)

        trX = trainX / 255.

        return trX, trY
    else:
        fd = open(os.path.join(path, 't10k-images-idx3-ubyte'))
        loaded = np.fromfile(file=fd, dtype=np.uint8)
        teX = loaded[16:].reshape((10000, 28, 28, 1)).astype(np.float)

        fd = open(os.path.join(path, 't10k-labels-idx1-ubyte'))
        loaded = np.fromfile(file=fd, dtype=np.uint8)
        teY = loaded[8:].reshape((10000)).astype(np.int32)

        teX = teX / 255.
        return teX, teY


teX, teY= load_mnist(is_training=False) # 获取数据

for j in range(400):
    fig, ax = plt.subplots(nrows=5, ncols=5, sharex='all', sharey='all', ) # 一张图片有5行5列个子图
    ax = ax.flatten()
    for i in range(25):
        img = teX[i+j*25].reshape(28, 28)
        ax[i].imshow(img, cmap='Greys', interpolation='nearest')
    ax[0].set_xticks([])
    ax[0].set_yticks([])
    plt.tight_layout()  # 自动紧凑布局
    plt.savefig(r"D:\test\%d.png" % j) # 要放在plt.show()前面，否则保存的是空白图片
    plt.show()