使用Tensorflow2.6, Keras实现MNIST分类

# TensorFlow and tf.keras
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
# Helper libraries
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
# import tensorflow_datasets as tfds

# 制作数据集
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255., x_test / 255.
# 要多加一维通道数才能训练
x_train = tf.expand_dims(x_train, -1)
x_test = tf.expand_dims(x_test, -1)
# 把标签转化成独热编码
y_train = np.float32(tf.keras.utils.to_categorical(y_train, num_classes=10))
y_test = np.float32(tf.keras.utils.to_categorical(y_test, num_classes=10))

# 超参数
batch_size = 64
epoch = 20

# 制作dataset，并且将训练集打乱
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size).shuffle(batch_size * 10)
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batc