快捷方式

嵌套IO函数

class torch.autograd.function.NestedIOFunction(*args, **kwargs)[源代码]

此类仅出于向后兼容性原因而存在。对于任何新的用例,请使用 Function 而不是此类。

backward(*gradients)[源代码]

共享反向实用程序。

返回类型

任何

backward_extended(*grad_output)[源代码]

用户定义的反向传播。

forward(*args)[源代码]

共享正向实用程序。

返回类型

任何

forward_extended(*input)[源代码]

用户定义的正向传播。

static jvp(ctx, *grad_inputs)

定义一个公式,使用前向模式自动微分对操作进行微分。

所有子类都必须覆盖此函数。它必须接受一个上下文 ctx 作为第一个参数,然后是与 forward() 相同数量的输入(对于 forward 函数的非张量输入,将传入 None),并且它应该返回与 forward() 相同数量的张量。每个参数都是相对于给定输入的梯度,每个返回值都应该是相对于相应输出的梯度。如果输出不是张量或函数相对于该输出不可微分,则可以对该输入传递 None 作为梯度。

您可以使用 ctx 对象将任何值从正向传播传递到此函数。

返回类型

任何

mark_dirty(*args, **kwargs)[源代码]

参见 Function.mark_dirty()

mark_non_differentiable(*args, **kwargs)[源代码]

参见 Function.mark_non_differentiable()

save_for_backward(*args)[source]

请参考 Function.save_for_backward().

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 的前向传递。

要么

  1. 用签名 forward(ctx, *args, **kwargs) 覆盖 forward。不覆盖 setup_context。为反向传播设置 ctx 发生在 forward 内部。

  2. 用签名 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,了解更多详细信息。

文档

获取 PyTorch 的全面开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并获取问题的答案

查看资源