如何将基于MNIST数据集的TensorFlow 1代码改写为TensorFlow 2兼容版本？

关键适配要点 #

TF1基于静态图机制，TF2默认启用动态图（Eager Execution），当兼容层方案无效时，建议直接改写为原生TF2代码，核心调整如下：

• 移除tf.Session()、tf.placeholder等TF1专属API，改用动态图直接运算或@tf.function装饰器构建静态图
• 用tf.keras.datasets.mnist加载数据集，替代TF1的input_data.read_data_sets接口
• 用tf.GradientTape或优化器内置方法处理梯度计算，替代TF1中依赖会话的optimizer.minimize()调用
• 严格对齐原代码的数据预处理逻辑、模型结构、超参数，确保输出结果一致

示例改写（原TF1代码→TF2代码） #

原TF1典型代码

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

# 加载数据
mnist = input_data.read_data_sets("MNIST_data/", one_hot=False)

# 占位符与模型参数
x = tf.placeholder(tf.float32, [None, 784])
y_ = tf.placeholder(tf.int32, [None])
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))

# 模型前向传播与损失计算
y = tf.matmul(x, W) + b
cross_entropy = tf.losses.sparse_softmax_cross_entropy(labels=y_, logits=y)
train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)

# 会话执行训练与评估
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for _ in range(1000):
        batch_xs, batch_ys = mnist.train.next_batch(100)
        sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})
    # 计算测试准确率
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.cast(y_, tf.int64))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    print(sess.run(accuracy, feed_dict={x: mnist.test.images, y_: mnist.test.labels}))

改写后的TF2代码

import tensorflow as tf

# 加载并预处理数据（与原代码逻辑对齐）
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype('float32') / 255.0
x_test = x_test.reshape(-1, 784).astype('float32') / 255.0

# 定义模型参数
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))

# 前向传播函数（用@tf.function模拟静态图行为）
@tf.function
def forward(x):
    return tf.matmul(x, W) + b

# 单步训练逻辑
@tf.function
def train_batch(x_batch, y_batch):
    with tf.GradientTape() as tape:
        logits = forward(x_batch)
        loss = tf.losses.sparse_softmax_cross_entropy(labels=y_batch, logits=logits)
    # 计算梯度并更新参数
    grads = tape.gradient(loss, [W, b])
    tf.keras.optimizers.SGD(learning_rate=0.5).apply_gradients(zip(grads, [W, b]))

# 训练循环
for _ in range(1000):
    # 随机抽取批量数据（对齐原代码的next_batch逻辑）
    batch_idx = tf.random.uniform(shape=[100], maxval=len(x_train), dtype=tf.int32)
    batch_xs = tf.gather(x_train, batch_idx)
    batch_ys = tf.gather(y_train, batch_idx)
    train_batch(batch_xs, batch_ys)

# 评估测试准确率
test_logits = forward(x_test)
correct = tf.equal(tf.argmax(test_logits, 1), tf.cast(y_test, tf.int64))
accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))
print(f"测试准确率: {accuracy.numpy():.4f}")

额外注意事项 #

• 若原代码包含自定义数据队列、复杂控制流，需用tf.data.Dataset替代TF1的数据队列API，用tf.cond、tf.while_loop（配合@tf.function）处理控制流逻辑
• 所有超参数（学习率、训练轮次、批量大小）必须与原代码完全一致，才能保证输出结果匹配
• 避免混合使用TF1兼容层与TF2原生API，否则易出现不可预期的兼容性问题

内容的提问来源于stack exchange，提问作者reh