你想了解除PyTorch外，其他主流深度学习框架实现MNIST手写数字分类的方法，我将以**TensorFlow 2.x（集成Keras高层API）**为例（目前最主流的PyTorch替代框架），实现与前文完全一致的功能，代码可直接运行，并附带详细解析，同时对比与PyTorch的核心差异。

一、环境准备

首先安装TensorFlow依赖：

# 安装CPU版本（通用）
pip install tensorflow
# 若有NVIDIA GPU，安装GPU版本（训练更快）
# pip install tensorflow-gpu

二、TensorFlow/Keras完整可运行代码

# 1. 导入核心库
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, losses, optimizers

# 2. 超参数设置（与PyTorch保持一致，方便对比）
BATCH_SIZE = 64
LEARNING_RATE = 0.001
EPOCHS = 5

# 3. 数据预处理与加载（TensorFlow内置MNIST数据集）
# 加载数据集（自动下载到本地缓存，无需手动指定路径）
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# 预处理：
# 1. 将图像数据转为float32类型，并归一化到[0,1]区间（替代PyTorch的标准化，简化实现，效果相近）
# 2. 无需手动展平，后续用Flatten层自动处理
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

# 调整数据形状：(样本数, 28, 28) → (样本数, 28, 28, 1)（适配Keras默认的输入格式）
x_train = tf.expand_dims(x_train, axis=-1)
x_test = tf.expand_dims(x_test, axis=-1)

# 构建数据加载器（类似PyTorch的DataLoader，支持批量加载和打乱）
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_ds = train_ds.shuffle(buffer_size=10000).batch(BATCH_SIZE)  # 打乱+批量
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_ds = test_ds.batch(BATCH_SIZE)  # 测试集无需打乱

# 4. 构建深度学习模型（与PyTorch结构完全一致的全连接网络）
# 使用Keras Sequential API（简洁直观，快速构建串行模型）
model = keras.Sequential([
    layers.Flatten(input_shape=(28, 28, 1)),  # 自动展平：(28,28,1) → 784，替代PyTorch的x.view()
    layers.Dense(128, activation="relu"),     # 全连接层+ReLU激活，对应PyTorch的nn.Linear+nn.ReLU
    layers.Dense(64, activation="relu"),      # 第二层全连接+ReLU
    layers.Dense(10)                          # 输出层（无激活，对应PyTorch的原始logits输出）
])

# 查看模型结构（可选，直观了解网络层级）
model.summary()

# 5. 定义损失函数与优化器（与PyTorch功能对应）
loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)  # 交叉熵损失（支持整数标签，内置Softmax）
optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)          # Adam优化器，与PyTorch一致

# 6. 定义训练指标（跟踪训练准确率，可选但实用）
train_loss_metric = keras.metrics.Mean(name='train_loss')  # 平均训练损失
train_acc_metric = keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')  # 训练准确率

# 7. 模型训练步骤（手动构建训练循环，贴近PyTorch逻辑，便于对比；也可直接用model.fit()简化实现）
@tf.function  # 装饰器：将Python函数转为TensorFlow计算图，加速训练
def train_step(x, y):
    # 梯度记录上下文（对应PyTorch的自动梯度计算，无需手动清空梯度，TensorFlow自动管理）
    with tf.GradientTape() as tape:
        logits = model(x, training=True)  # 前向传播，training=True启用训练模式
        loss = loss_fn(y, logits)         # 计算损失
    
    # 反向传播+参数更新（对应PyTorch的loss.backward() + optimizer.step()）
    gradients = tape.gradient(loss, model.trainable_variables)  # 计算梯度
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))  # 更新参数

    # 更新训练指标
    train_loss_metric.update_state(loss)
    train_acc_metric.update_state(y, logits)

# 8. 模型评估步骤（对应PyTorch的evaluate函数）
@tf.function
def test_step(x, y):
    logits = model(x, training=False)  # 前向传播，training=False启用评估模式
    loss = loss_fn(y, logits)          # 计算测试损失

    # 更新测试指标
    test_loss_metric.update_state(loss)
    test_acc_metric.update_state(y, logits)

# 9. 启动训练与评估
# 初始化测试指标
test_loss_metric = keras.metrics.Mean(name='test_loss')
test_acc_metric = keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

for epoch in range(EPOCHS):
    # 重置训练指标
    train_loss_metric.reset_states()
    train_acc_metric.reset_states()

    # 遍历训练集
    for batch_idx, (x_batch, y_batch) in enumerate(train_ds):
        train_step(x_batch, y_batch)

        # 每100个批次打印训练信息（与PyTorch一致）
        if batch_idx % 100 == 99:
            print(f'Epoch [{epoch+1}/{EPOCHS}], Batch [{batch_idx+1}/{len(train_ds)}], '
                  f'Loss: {train_loss_metric.result().numpy():.4f}, '
                  f'Accuracy: {train_acc_metric.result().numpy() * 100:.2f}%')

    # 打印本轮训练整体指标
    print(f'\nEpoch [{epoch+1}/{EPOCHS}] Train Summary - Loss: {train_loss_metric.result().numpy():.4f}, '
          f'Accuracy: {train_acc_metric.result().numpy() * 100:.2f}%\n')

    # 重置测试指标
    test_loss_metric.reset_states()
    test_acc_metric.reset_states()

    # 遍历测试集评估
    for x_batch, y_batch in test_ds:
        test_step(x_batch, y_batch)

    # 打印测试指标
    print(f'Test Summary - Loss: {test_loss_metric.result().numpy():.4f}, '
          f'Accuracy: {test_acc_metric.result().numpy() * 100:.2f}%\n')

# 10. 保存模型（两种常用格式）
# 格式1：Keras完整模型（含结构、权重、优化器状态，可直接加载训练）
model.save('mnist_keras_full.h5')
# 格式2：仅保存权重（与PyTorch的state_dict对应）
model.save_weights('mnist_keras_weights.h5')
print("模型已保存（完整模型：mnist_keras_full.h5，权重：mnist_keras_weights.h5）")

# 加载模型示例（可选）
# loaded_full_model = keras.models.load_model('mnist_keras_full.h5')
# loaded_weight_model = keras.Sequential(...)  # 先构建相同结构
# loaded_weight_model.load_weights('mnist_keras_weights.h5')

三、核心模块解析（对比PyTorch）

1. 数据处理差异

2. 模型构建差异

• TensorFlow/Keras提供两种核心建模方式：
  1. Sequential API（上文使用）：简洁高效，适用于串行网络结构（如本次全连接网络），无需手动定义forward方法。
  2. Functional API：灵活强大，适用于多输入、多输出、残差连接等复杂网络。
• 与PyTorch对比：Keras无需继承nn.Module，无需手动实现forward方法，层级定义更简洁；PyTorch则更灵活，支持自定义复杂传播逻辑。

3. 训练流程核心差异

这是两个框架的核心区别，需重点掌握：

4. 损失函数与优化器

• 功能完全对应：SparseCategoricalCrossentropy（TensorFlow）= nn.CrossEntropyLoss（PyTorch），均支持整数标签，内置Softmax；Adam优化器参数和效果一致。
• 差异：TensorFlow的损失函数和优化器需单独实例化，PyTorch则直接传入模型参数初始化优化器。

5. 模型保存与加载

• TensorFlow两种保存格式：
  1. 完整模型（.h5）：含网络结构、权重、优化器状态，可直接加载继续训练或推理，无需重新构建模型。
  2. 仅权重（.h5）：与PyTorch的state_dict对应，需先构建相同结构的模型，再加载权重。
• PyTorch仅支持权重（state_dict）保存，需手动构建模型后加载权重。

四、简化实现：使用model.fit()（Keras特色）

Keras的最大优势是高度封装，可省略手动训练循环，用model.fit()一行实现训练，代码更简洁：

# 省略数据加载与预处理（同上文）
# 省略模型构建（同上文）

# 编译模型（绑定损失函数、优化器、评估指标）
model.compile(
    optimizer=optimizers.Adam(learning_rate=LEARNING_RATE),
    loss=losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

# 直接训练（自动处理批次、训练/评估切换、打印日志）
history = model.fit(
    x_train, y_train,
    batch_size=BATCH_SIZE,
    epochs=EPOCHS,
    validation_data=(x_test, y_test)  # 自动评估测试集
)

# 评估模型
test_loss, test_acc = model.evaluate(x_test, y_test, batch_size=BATCH_SIZE)
print(f'最终测试准确率：{test_acc*100:.2f}%')

五、其他框架补充说明

1. MXNet：与PyTorch语法相近，核心流程一致，数据集加载需通过mxnet.gluon.data.vision.MNIST，模型继承gluon.Block，训练逻辑与PyTorch类似。
2. PaddlePaddle（百度）：国产框架，文档中文友好，API设计参考TensorFlow和PyTorch，MNIST分类实现可参考其官方教程，核心流程（数据-模型-损失-训练-评估）与主流框架一致。
3. 核心共性：所有深度学习框架实现分类任务的核心逻辑均为「数据预处理→模型构建→损失计算→反向传播→参数更新→评估」，差异仅在于API设计和细节实现。

六、总结

1. TensorFlow/Keras是PyTorch的主流替代框架，Sequential API简洁高效，model.fit()可快速实现训练，手动循环则更贴近PyTorch逻辑，便于迁移学习。
2. 两个框架核心功能一一对应，差异主要在梯度管理、模型构建和训练循环的实现形式。
3. 本次代码与PyTorch实现的模型结构、超参数完全一致，测试准确率同样可达97%以上，模型可直接保存和加载。
4. 其他框架（MXNet/PaddlePaddle）核心流程一致，仅需适配对应API即可实现相同功能。