torch.autograd.function.FunctionCtx.save_for_backward
- FunctionCtx.save_for_backward(*tensors)[源代码]
-
为将来调用
backward()
保存给定的张量。save_for_backward
应该在setup_context()
或forward()
方法中最多调用一次,并且只允许与张量一起使用。所有计划在反向传递中使用的张量应使用
save_for_backward
(而不是直接保存到ctx
)来保存,以防止计算不正确的梯度和内存泄漏,并启用已保存张量挂钩的应用。参见torch.autograd.graph.saved_tensors_hooks
。注意,如果你保存了既不是
forward()
输入也不是输出的中间张量用于反向传播,你的自定义函数可能不支持双反向传播。对于不支持双反向传播的自定义函数,应该用@once_differentiable
装饰其backward()
方法,以便在执行双反向传播时引发错误。如果你想支持双反向传播,可以在反向传播期间根据输入重新计算中间张量,或者将这些中间张量作为自定义函数的输出返回。更多详情请参阅双反向传播教程。在
backward()
中,可以通过saved_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)