首先直接给结论：TensorFlow 的 Session 并没有内置的 sess.undo_last() 这类一键回滚方法。原因很简单：当你运行训练操作（train_op）时，模型的可训练变量（权重、偏置等）是直接在内存中被修改的，Session 不会自动记录每一步操作的状态快照，所以没法直接撤销单步训练。

不过针对你提到的「损失值为NaN时需要回滚」的场景，我们有几个可靠的替代方案，同时也能适配多GPU训练：

方案1：用检查点（Checkpoint）提前保存正常状态 #

这是最常用也最稳妥的方法——在训练过程中定期保存模型的正常状态，当遇到NaN时，加载上一个正常的检查点即可回滚。

示例代码贴合你的需求：

import tensorflow as tf
import numpy as np

# 假设你已经定义了模型、train_op、loss等
saver = tf.train.Saver()
last_normal_step = 0
# 先保存初始状态
saver.save(sess, './checkpoints/normal_model', global_step=last_normal_step)

for step in range(num_epoch):
    _, loss_value = sess.run([train_op, loss])
    if np.isnan(loss_value):
        # 加载上一个正常的检查点
        saver.restore(sess, './checkpoints/normal_model-{}'.format(last_normal_step))
        print(f"Step {step} encountered NaN loss, rolled back to step {last_normal_step}")
        break
    # 如果当前步骤正常，更新记录并保存
    last_normal_step = step
    saver.save(sess, './checkpoints/normal_model', global_step=last_normal_step)

你可以根据需求调整保存频率——比如每10步保存一次，或者只在验证集损失下降时保存，减少IO开销。

方案2：手动保存变量副本（适合小模型） #

如果你的模型规模不大，可以在每步训练前手动复制当前所有可训练变量的值，遇到NaN时再把变量赋值回之前的状态：

trainable_vars = tf.trainable_variables()
# 先保存初始状态
current_var_values = sess.run(trainable_vars)

for step in range(num_epoch):
    # 训练前保存当前状态
    prev_var_values = current_var_values
    _, loss_value = sess.run([train_op, loss])
    if np.isnan(loss_value):
        # 回滚变量到训练前的状态
        for var, val in zip(trainable_vars, prev_var_values):
            sess.run(var.assign(val))
        print(f"Step {step} got NaN, rolled back to previous state")
        break
    # 更新当前状态为正常训练后的状态
    current_var_values = sess.run(trainable_vars)

这种方法不需要磁盘IO，但会占用额外内存，大模型慎用。

多GPU训练场景的适配 #

不管是用检查点还是变量副本，核心思路和单GPU一致，但要注意几个细节：

• 检查点方法：如果你用的是TensorFlow的分布式策略（比如tf.distribute.MirroredStrategy），tf.train.Checkpoint或者Keras的model.save_weights()已经原生支持分布式场景——只要在保存和加载时处于策略的上下文管理器中，就能正确同步所有GPU的变量状态，回滚后所有GPU都会使用恢复后的全局变量。
• 变量副本方法：多GPU训练时，通常会有一个主变量集合（比如放在CPU），各个GPU的变量是主变量的副本。你只需要复制主变量的值即可，回滚主变量后，下次训练时GPU会自动同步最新的主变量状态，不需要单独处理每个GPU的副本。

总结一下：虽然没有原生的撤销方法，但通过提前保存状态的方式完全可以实现你要的需求，而且这些方案在多GPU场景下也能很好地工作。

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