评价此页

torch.autograd.function.FunctionCtx.save_for_backward#

FunctionCtx.save_for_backward(*tensors)[source]#

为未来的 backward() 调用保存给定的张量。

save_for_backward 最多只能调用一次,可以在 setup_context()forward() 方法中调用,并且只能使用 tensors。

所有打算在 backward 传播中使用但不是 forward 函数的输入或输出的 tensor 都应使用 save_for_backward 保存(而不是直接保存在 ctx 上),以防止梯度不正确和内存泄漏,并启用已保存 tensor hook 的应用。请参阅 torch.autograd.graph.saved_tensors_hooks。有关更多详细信息,请参阅 扩展 torch.autograd

请注意,如果保存了中间 tensors(即既不是 forward() 的输入也不是输出的 tensors),那么自定义 Function 可能不支持二阶反向传播。不支持二阶反向传播的自定义 Function 应该用 @once_differentiable 装饰其 backward() 方法,这样执行二阶反向传播时会报错。如果您希望支持二阶反向传播,可以根据反向传播过程中的输入重新计算中间值,或者将中间值作为自定义 Function 的输出返回。有关更多详细信息,请参阅 二阶反向传播教程

backward() 中,可以通过 saved_tensors 属性访问已保存的 tensors。在将它们返回给用户之前,会进行检查以确保它们没有被用于任何就地修改其内容的运算。

参数也可以是 None。这不会执行任何操作。

有关如何使用此方法的更多详细信息,请参阅 扩展 torch.autograd

示例

>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x: torch.Tensor, y: torch.Tensor, z: int):
>>>         w = x * z
>>>         out = x * y + y * z + w * y
>>>         ctx.save_for_backward(x, y, w, out)
>>>         ctx.z = z  # z is not a tensor
>>>         return out
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, grad_out):
>>>         x, y, w, out = ctx.saved_tensors
>>>         z = ctx.z
>>>         gx = grad_out * (y + y * z)
>>>         gy = grad_out * (x + z + w)
>>>         gz = None
>>>         return gx, gy, gz
>>>
>>> a = torch.tensor(1., requires_grad=True, dtype=torch.double)
>>> b = torch.tensor(2., requires_grad=True, dtype=torch.double)
>>> c = 4
>>> d = Func.apply(a, b, c)