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)
本页目录