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()
的输入也不是输出的张量)被保存用于反向传播,您的自定义 Function 可能不支持二次反向传播。不支持二次反向传播的自定义 Function 应该使用@once_differentiable
装饰其backward()
方法,以便在执行二次反向传播时引发错误。如果您想支持二次反向传播,可以在反向传播期间根据输入重新计算中间张量,或者将中间张量作为自定义 Function 的输出返回。更多详细信息请参阅二次反向传播教程。在
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)