用 TensorFlow 2 与 Keras 训练 MNIST 手写数字分类模型

TensorFlow 2 把 Keras 作为默认的高层建模接口，让「加载数据 → 搭模型 → 编译 → 训练 → 评估」这条链路可以用很少的代码走完。本文要完成的具体任务是：用 TensorFlow 2 和 Keras 训练一个识别手写数字（0–9）的神经网络，并在测试集上验证它的准确率。适合已经会写一点 Python、想第一次把模型真正跑起来的人。整个流程在浏览器里的 Notebook 环境（例如 Google Colab）就能完成，不必先在本地配好 GPU 环境；如果你更习惯本地开发，也可以按官方安装说明装好 TensorFlow 2 再跟着做。想先了解这个框架的整体定位，可以看站内的 TensorFlow 条目。

这篇教程覆盖的内容包括：导入并确认 TensorFlow 版本、加载 MNIST 并做像素归一化、用 Sequential 堆叠 Flatten / Dense / Dropout 层、理解模型输出的是 logits 而不是概率、用 SparseCategoricalCrossentropy 定义损失、用 model.fit 训练 5 个 epoch、用 model.evaluate 评估，最后给出一个从头到尾可运行的最小示例和若干容易踩的坑。

准备工作

运行环境

最省事的方式是使用浏览器内的 Notebook 环境。这类环境已经预装好 TensorFlow，打开就能运行 Python 代码，不需要自己处理 CUDA、cuDNN 之类的依赖。使用步骤是：

• 打开 Notebook 页面后，在菜单栏右上角选择连接（CONNECT），连上一个 Python 运行时。
• 想一次性跑完全部代码，选择「Runtime > Run all」。
• 想逐格运行，把鼠标悬停在某个代码单元上，点击运行该单元的图标。

如果你要在自己的开发环境里跑，需要先安装 TensorFlow 2 包。安装前请确认 pip 已升级到较新版本，否则可能装不上或装到旧版本。具体的安装命令、Python 版本要求与 GPU 支持情况会随版本变化，请以 TensorFlow 官网当前的安装说明为准。

需要的前置知识

• 能读懂基本的 Python 语法，知道列表、函数、导入模块。
• 知道什么是「张量」的大致概念即可，不要求会手写反向传播。
• 不需要提前准备数据集文件，MNIST 由 Keras 内置接口直接下载。

确认版本

第一步永远是确认环境里的 TensorFlow 版本，避免照着新版本的写法在旧版本上跑。导入并打印版本号：

import tensorflow as tf
print("TensorFlow version:", tf.__version__)

在带 GPU 的环境里，这段代码通常会顺带打印一批日志，例如提示 cuFFT、cuDNN、cuBLAS 的工厂函数被重复注册，或者提示 computation placer already registered。这些是底层库在初始化时的重复链接提示，不影响后续计算，可以忽略。真正需要关注的是最后打印出来的版本号。

操作步骤

第一步：加载 MNIST 数据集

MNIST 是一组 28×28 像素的手写数字灰度图，训练集和测试集由 tf.keras.datasets.mnist 直接提供。加载后得到四个数组：训练图像、训练标签、测试图像、测试标签。

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()

原始像素值的范围是 0 到 255。神经网络对输入尺度很敏感，所以要把它们缩放到 0 到 1 之间，做法是整体除以 255.0：

x_train, x_test = x_train / 255.0, x_test / 255.0

这一步同时把样本数据从整数变成了浮点数。别小看这一行，如果跳过归一化，训练初期的损失会非常大，收敛也会明显变慢。

第二步：构建模型

用 tf.keras.Sequential 按顺序堆叠层。这个模型包含四层：

• Flatten：把 28×28 的二维图像拉平成一维向量，共 784 个元素。
• Dense(128, activation='relu')：128 个神经元的全连接层，使用 ReLU 激活。
• Dropout(0.2)：训练时随机丢弃 20% 的神经元，用来抑制过拟合。
• Dense(10)：输出 10 个值，对应 0–9 十个类别，不加激活函数。

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
])

运行这段代码时，Keras 可能给出一个 UserWarning，提示不要在层上直接传 input_shape / input_dim，而建议在 Sequential 模型的第一层使用 Input(shape) 对象。按提示改写更规范：

model = tf.keras.models.Sequential([
    tf.keras.layers.Input(shape=(28, 28)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
])

Sequential 适合「每层一个输入张量、一个输出张量」的线性堆叠场景。层本身是带有可训练变量的函数，绝大多数 TensorFlow 模型都是由层组合而成的。

第三步：理解模型输出的是 logits

最后一层没有激活函数，所以模型输出的是 logits（对数几率），每个类别一个分数，而不是概率。取一条训练样本看一下：

predictions = model(x_train[:1]).numpy()
predictions

输出形如一个 1×10 的数组，里面是正负不一的浮点数。要把这些分数变成各类别的概率，用 tf.nn.softmax：

tf.nn.softmax(predictions).numpy()

转换后每个值都在 0 到 1 之间，且十个值加起来为 1。这里有一个值得记住的设计取舍：可以把 softmax 直接写进最后一层的激活函数里，这样模型输出更直观；但不推荐这么做，因为一旦输出层带了 softmax，就很难为所有模型提供精确且数值稳定的损失计算。所以惯例是让模型输出 logits，把 softmax 留给推理阶段或损失函数内部处理。

第四步：定义损失函数

多分类、标签是整数（而不是 one-hot 向量）的场景，用 SparseCategoricalCrossentropy，并设置 from_logits=True 告诉它输入是 logits：

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

这个损失函数接收真实标签向量和 logits 向量，为每个样本返回一个标量损失。它的含义是「真实类别对应概率的负对数」：模型越确信正确答案，损失越接近 0。

可以用一条样本验证一下量级。未训练的模型输出接近随机，十个类别各约 1/10，所以初始损失应该接近 -tf.math.log(1/10)，也就是 2.3 左右：

loss_fn(y_train[:1], predictions).numpy()

如果得到的结果在 2 到 3 之间，说明损失函数的配置是对的。如果明显偏离这个量级，通常意味着 from_logits 设错了，或者模型输出层被加了 softmax。

第五步：编译模型

训练之前用 Model.compile 配置优化器、损失和评估指标：

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

• optimizer='adam'：使用 Adam 优化器，按字符串名称调用即可。
• loss=loss_fn：传入上一步定义的损失函数对象。
• metrics=['accuracy']：训练和评估过程中额外报告准确率。

第六步：训练模型

调用 Model.fit 让模型在训练数据上迭代，通过优化器不断调整参数以最小化损失：

model.fit(x_train, y_train, epochs=5)

epochs=5 表示把整个训练集过 5 遍。每个 epoch 结束时，进度条后面会显示当前的 loss 和 accuracy。训练过程中同样会刷出大量底层日志，例如 XLA 服务初始化、检测到的 GPU 设备名称与显存大小、计算能力版本等，这些都属于正常输出。

第七步：评估模型

训练完成后，在测试集上评估泛化表现：

model.evaluate(x_test, y_test, verbose=2)

返回的第一个值是测试集上的损失，第二个值是准确率。因为测试数据没有参与训练，这个准确率比训练过程中报告的数字更能说明模型的真实水平。

第八步（可选）：让模型直接输出概率

如果希望模型直接给出概率而不是 logits，可以再包一层 softmax，得到一个可以单独使用的新模型：

probability_model = tf.keras.Sequential([
    model,
    tf.keras.layers.Softmax()
])

probability_model(x_test[:5])

这样对前 5 条测试样本调用时，得到的就是每类概率。注意这只是推理侧的包装，原模型的训练方式不受影响。

一个完整示例

把上面的步骤串起来，下面这段代码从导入到评估一次跑完，可以直接粘贴进 Notebook 的单个单元执行：

import tensorflow as tf

print("TensorFlow version:", tf.__version__)

# 1. 加载并归一化数据
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# 2. 构建模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Input(shape=(28, 28)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
])

# 3. 定义损失
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 4. 编译
model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])

# 5. 训练
model.fit(x_train, y_train, epochs=5)

# 6. 评估
model.evaluate(x_test, y_test, verbose=2)

执行顺序上有两点值得注意：model.compile 必须在 model.fit 之前调用，否则会报错；loss_fn 要在编译之前定义好，因为编译时就要把它绑定到模型上。如果中途修改了模型结构（比如增删层），需要重新编译，否则优化器和损失配置不会自动更新。

跑完之后，你会得到测试集上的损失与准确率两个数字。想进一步观察模型对具体样本的预测，可以用 model.predict(x_test[:5]) 取前 5 条样本的 logits，再用 tf.nn.softmax 转成概率，与 y_test[:5] 里的真实标签对照。

注意事项

日志与警告不等于出错

在 GPU 环境中，导入 TensorFlow 和构建模型时经常出现成片的 E 与 W 级别日志，例如 cuFFT / cuDNN / cuBLAS 工厂重复注册、computation placer already registered、XLA 服务初始化信息、检测到的 GPU 型号与显存大小等。这些是底层运行时的正常提示，不代表训练失败。判断是否真的出错，看 Python 是否抛出异常、以及 fit 与 evaluate 是否正常返回指标。

input_shape 的弃用提示

在 Flatten 层上传 input_shape 会触发 UserWarning。虽然代码仍能运行，但建议改成在模型第一层使用 Input(shape=(28, 28))，这样更符合当前 Keras 的推荐写法，也能避免后续版本升级时出问题。

不要把 softmax 塞进输出层

输出层保持线性、由损失函数内部处理 softmax，是数值稳定性更好的做法。如果为了「输出好看」而在最后一层加 softmax，同时损失函数又设了 from_logits=True，就会造成重复变换，损失量级和梯度都会失真。

归一化不能省

像素值不除以 255.0 时，输入尺度在 0–255，配合 Adam 和交叉熵虽然有时也能训练，但收敛更慢、初始损失更大。把输入缩放到 0–1 是这个任务的标准做法。

版本与安装信息以官网为准

TensorFlow 的版本号、Python 版本要求、GPU 驱动与 CUDA 兼容矩阵、各平台安装方式都在持续变化，本文不锁定具体版本。安装前请查阅 TensorFlow 官网当前的安装说明；如果使用托管 Notebook 环境，其预装版本同样以该环境当时的实际输出为准，用 tf.__version__ 打印确认即可。

训练轮数与过拟合

示例中固定使用 5 个 epoch。轮数太少模型欠拟合，轮数太多则可能在训练集上过拟合、测试集准确率反而下降。判断依据是同时观察训练过程中的 loss / accuracy 与 evaluate 返回的测试指标，两者差距明显拉大时就该考虑减少轮数或加强正则化（例如调高 Dropout 的比例）。