PyTorch中@torch.no_grad的TensorFlow等效实现及使用疑问

TensorFlow提供了多种和PyTorch中@torch.no_grad等价的方式，针对不同场景可灵活选择：

1. 最接近PyTorch风格：tf.no_grad() 上下文/装饰器 #

和PyTorch的@torch.no_grad用法几乎一致，既可以作为上下文管理器包裹代码块，也可以作为装饰器修饰函数，直接禁用梯度计算：

作为上下文管理器 #

x = tf.Variable([1.0, 2.0])
with tf.no_grad():
    # 此代码块内的所有操作都不会追踪梯度
    y = x * 3
    print(y)  # 输出 tf.Tensor([3. 6.], shape=(2,), dtype=float32)
    # 尝试计算梯度会返回None
    grad = tf.gradients(y, x)
    print(grad)  # 输出 None

作为装饰器 #

@tf.no_grad()
def predict(model, input_data):
    # 函数内的模型推理不会计算梯度
    return model(input_data)

# 使用时直接调用
output = predict(my_model, test_input)

2. 局部禁用梯度：tf.stop_gradient() #

如果只需要对某一个张量操作禁用梯度，而非整个代码块，可用tf.stop_gradient()包裹该操作的结果，切断该张量的梯度传播路径：

x = tf.Variable(2.0)
with tf.GradientTape() as tape:
    y = x + 1
    # 对y*2的结果禁用梯度，梯度只会传到y为止
    z = tf.stop_gradient(y * 2)
    loss = z + 3

grad = tape.gradient(loss, x)
print(grad)  # 输出 tf.Tensor(1.0, shape=(), dtype=float32)
# 因为z的梯度被切断，所以loss对x的梯度等于y对x的梯度（即1）

3. GradientTape内临时停止记录：tape.stop_recording() #

你提到的tape.stop_recording()针对已开启tf.GradientTape的场景，用来临时停止梯度记录，不需要把tape传入模型，只需在tape的上下文内嵌套该上下文管理器即可：

x1 = tf.Variable([1.0])
x2 = tf.Variable([2.0])
model = tf.keras.layers.Dense(1)

with tf.GradientTape() as tape:
    # 这段操作会被tape记录梯度
    pred1 = model(x1)
    loss1 = tf.square(pred1 - 3)
    
    # 临时停止梯度记录
    with tape.stop_recording():
        # 这段操作不会被tape记录，即使是模型的推理
        pred2 = model(x2)
        loss2 = tf.square(pred2 - 4)
    
    # 退出stop_recording上下文后自动恢复记录
    total_loss = loss1 + loss2

# 计算梯度时，只有loss1相关的操作会被考虑，loss2的梯度不会被计算
grads = tape.gradient(total_loss, model.trainable_variables)

这种方式适合在同一个GradientTape上下文里，部分操作需要记录梯度、部分不需要的场景。

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