嵌套IO函数¶
- class torch.autograd.function.NestedIOFunction(*args, **kwargs)[源代码]¶
此类仅出于向后兼容性原因而存在。对于任何新的用例,请使用
Function
而不是此类。- static jvp(ctx, *grad_inputs)¶
定义一个公式,使用前向模式自动微分对操作进行微分。
所有子类都必须覆盖此函数。它必须接受一个上下文
ctx
作为第一个参数,然后是与forward()
相同数量的输入(对于 forward 函数的非张量输入,将传入 None),并且它应该返回与forward()
相同数量的张量。每个参数都是相对于给定输入的梯度,每个返回值都应该是相对于相应输出的梯度。如果输出不是张量或函数相对于该输出不可微分,则可以对该输入传递 None 作为梯度。您可以使用
ctx
对象将任何值从正向传播传递到此函数。- 返回类型
- save_for_forward(*tensors)¶
将给定的张量保存起来,以便将来调用
jvp()
时使用。save_for_forward
应该最多调用一次,可以在setup_context()
或forward()
方法中调用,所有参数都应该是张量。在
jvp()
中,可以通过saved_tensors
属性访问保存的对象。参数也可以是
None
。这将是一个空操作。请参考 扩展 torch.autograd,了解有关如何使用此方法的更多详细信息。
- 示例:
>>> class Func(torch.autograd.Function): >>> @staticmethod >>> def forward(ctx, x: torch.Tensor, y: torch.Tensor, z: int): >>> ctx.save_for_backward(x, y) >>> ctx.save_for_forward(x, y) >>> ctx.z = z >>> return x * y * z >>> >>> @staticmethod >>> def jvp(ctx, x_t, y_t, _): >>> x, y = ctx.saved_tensors >>> z = ctx.z >>> return z * (y * x_t + x * y_t) >>> >>> @staticmethod >>> def vjp(ctx, grad_out): >>> x, y = ctx.saved_tensors >>> z = ctx.z >>> return z * grad_out * y, z * grad_out * x, None >>> >>> a = torch.tensor(1., requires_grad=True, dtype=torch.double) >>> t = torch.tensor(1., dtype=torch.double) >>> b = torch.tensor(2., requires_grad=True, dtype=torch.double) >>> c = 4 >>> >>> with fwAD.dual_level(): >>> a_dual = fwAD.make_dual(a, t) >>> d = Func.apply(a_dual, b, c)
- property saved_tensors¶
请参考
Function.saved_tensors()
.
- set_materialize_grads(value)¶
设置是否物化梯度张量。默认值为
True
。这应该只在
setup_context()
或forward()
方法中调用。如果为
True
,则在调用backward()
和jvp()
方法之前,未定义的梯度张量将扩展为全零张量。- 示例:
>>> class SimpleFunc(Function): >>> @staticmethod >>> def forward(ctx, x): >>> return x.clone(), x.clone() >>> >>> @staticmethod >>> @once_differentiable >>> def backward(ctx, g1, g2): >>> return g1 + g2 # No check for None necessary >>> >>> # We modify SimpleFunc to handle non-materialized grad outputs >>> class Func(Function): >>> @staticmethod >>> def forward(ctx, x): >>> ctx.set_materialize_grads(False) >>> ctx.save_for_backward(x) >>> return x.clone(), x.clone() >>> >>> @staticmethod >>> @once_differentiable >>> def backward(ctx, g1, g2): >>> x, = ctx.saved_tensors >>> grad_input = torch.zeros_like(x) >>> if g1 is not None: # We must check for None now >>> grad_input += g1 >>> if g2 is not None: >>> grad_input += g2 >>> return grad_input >>> >>> a = torch.tensor(1., requires_grad=True) >>> b, _ = Func.apply(a) # induces g2 to be undefined
- static setup_context(ctx, inputs, output)¶
有两种方法可以定义 autograd.Function 的前向传递。
要么
用签名
forward(ctx, *args, **kwargs)
覆盖 forward。不覆盖setup_context
。为反向传播设置 ctx 发生在forward
内部。用签名
forward(*args, **kwargs)
覆盖 forward 并覆盖setup_context
。为反向传播设置 ctx 发生在setup_context
内部(而不是在forward
内部)。
请参考
torch.autograd.Function.forward()
和 扩展 torch.autograd,了解更多详细信息。- 返回类型
- static vjp(ctx, *grad_outputs)¶
定义使用反向模式自动微分来对操作进行微分的公式。
所有子类都必须覆盖此函数。(定义此函数等效于定义
vjp
函数。)它必须接受上下文
ctx
作为第一个参数,后面跟着与forward()
返回的输出一样多的输出(如果 forward 函数的非张量输出,则传递 None),并且它应该返回与forward()
的输入一样多的张量。每个参数都是相对于给定输出的梯度,每个返回值都应该是相对于相应输入的梯度。如果输入不是张量,或者是一个不需要计算梯度的张量,那么可以将 None 作为该输入的梯度传入。上下文可以用来检索在正向传递过程中保存的张量。它还有一个属性
ctx.needs_input_grad
,它是一个布尔值元组,表示每个输入是否需要计算梯度。例如,backward()
将会有ctx.needs_input_grad[0] = True
,如果forward()
的第一个输入需要计算相对于输出的梯度。- 返回类型
- static vmap(info, in_dims, *args)¶
定义此 autograd.Function 在
torch.vmap()
下的行为。为了让
torch.autograd.Function()
支持torch.vmap()
,你必须要么覆盖这个静态方法,要么将generate_vmap_rule
设置为True
(你不能同时做这两件事)。如果你选择覆盖这个静态方法:它必须接受
一个
info
对象作为第一个参数。info.batch_size
指定正在进行 vmap 操作的维度的尺寸,而info.randomness
是传递给torch.vmap()
的随机性选项。一个
in_dims
元组作为第二个参数。对于args
中的每个参数,in_dims
都有一个对应的Optional[int]
。如果参数不是张量,或者参数没有进行 vmap 操作,则为None
,否则,它是一个整数,指定正在进行 vmap 操作的张量的哪个维度。*args
,它与forward()
的参数相同。
vmap 静态方法的返回值是
(output, out_dims)
的元组。与in_dims
相似,out_dims
应该与output
的结构相同,并为每个输出包含一个out_dim
,指定输出是否具有 vmap 维度,以及它在哪个索引处。请参考 使用 autograd.Function 扩展 torch.func,了解更多详细信息。