InplaceFunction

class torch.autograd.function.InplaceFunction(inplace=False)

此类仅出于向后兼容的原因而存在。 对于任何新的用例，请使用 Function 代替此类。

static backward(ctx, *grad_outputs)

定义一个用于反向模式自动微分的操作的微分公式。

此函数应由所有子类重写。 （定义此函数等同于定义 vjp 函数。）

它必须接受一个上下文 ctx 作为第一个参数，后面跟着与 forward() 返回的输出一样多的参数（对于 forward 函数中的非张量输出，将传递 None），并且它应该返回与 forward() 的输入一样多的张量。每个参数是相对于给定输出的梯度，每个返回值应该是相对于相应输入的梯度。如果输入不是张量或是不需要梯度的张量，你可以为该输入传递 None 作为梯度。

上下文可以用于检索在前向传递过程中保存的张量。它还具有一个属性 ctx.needs_input_grad，作为表示每个输入是否需要梯度的布尔元组。例如， backward() 将具有 ctx.needs_input_grad[0] = True，如果 forward() 的第一个输入需要相对于输出的梯度计算。

Return type

任意

static forward(ctx, *args, **kwargs)

定义自定义 autograd 函数的 forward 方法。

此函数应由所有子类重写。 定义forward有两种方式：

用法 1（组合正向和上下文）：

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• 它必须接受一个上下文 ctx 作为第一个参数，后面可以跟随任意数量的参数（张量或其他类型）。

• 参见 组合或单独的 forward() 和 setup_context() 了解更多详情

用法 2（分离的前向和上下文）：

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• forward 不再接受 ctx 参数。

• 相反，您还必须重写 torch.autograd.Function.setup_context() 静态方法以处理设置 ctx 对象。 output 是前向传播的输出，inputs 是前向传播输入的元组。

• 参见 扩展 torch.autograd 了解更多详情

上下文可以用于存储任意数据，这些数据可以在反向传播过程中被检索。张量不应直接存储在ctx上（尽管目前出于向后兼容性考虑，这一点并未强制执行）。相反，如果张量旨在用于backward（等效地，vjp），则应使用ctx.save_for_backward()保存；如果张量旨在用于jvp，则应使用ctx.save_for_forward()保存。

Return type

任意

static jvp(ctx, *grad_inputs)

定义一个用于前向模式自动微分的操作的公式。

此函数应由所有子类重写。 它必须接受上下文 ctx 作为第一个参数，后面跟随与 forward() 接收的输入数量相同的输入（对于 forward 函数的非张量输入，将传递 None）， 并且它应返回与 forward() 输出数量相同的张量。每个参数是相对于给定输入的梯度， 并且每个返回值应是相对于相应输出的梯度。如果某个输出不是张量或函数对该输出不可微，您可以为该输入传递 None 作为梯度。

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

Return type

任意

mark_dirty(*args)

将给定的张量标记为在原地操作中被修改。

这最多只能调用一次，只能在 forward() 方法内部调用，并且所有参数都应该是输入。

在调用forward() 时，每个被就地修改的张量都应传递给此函数，以确保我们检查的正确性。 无论函数是在修改之前还是之后调用，都无关紧要。

Examples::

>>> class Inplace(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         x_npy = x.numpy() # x_npy 与 x 共享存储
>>>         x_npy += 1
>>>         ctx.mark_dirty(x)
>>>         return x
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, grad_output):
>>>         return grad_output
>>>
>>> a = torch.tensor(1., requires_grad=True, dtype=torch.double).clone()
>>> b = a * a
>>> Inplace.apply(a)  # 这会导致错误的梯度！
>>>                   # 但引擎不会知道，除非我们 mark_dirty
>>> b.backward() # RuntimeError: 梯度计算所需的变量之一已被就地操作修改
>>>              # 计算已被就地操作修改

mark_non_differentiable(*args)

将输出标记为不可微分。

这最多只能调用一次，只能在 forward() 方法内部调用，并且所有参数都应该是张量输出。

这将标记输出为不需要梯度，从而提高反向传播计算的效率。你仍然需要在 backward() 中接受每个输出的梯度，但它总是一个与相应输出形状相同的零张量。

This is used e.g. for indices returned from a sort. See example::

>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         sorted, idx = x.sort()
>>>         ctx.mark_non_differentiable(idx)
>>>         ctx.save_for_backward(x, idx)
>>>         return sorted, idx
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):  # 仍然需要接受 g2
>>>         x, idx = ctx.saved_tensors
>>>         grad_input = torch.zeros_like(x)
>>>         grad_input.index_add_(0, idx, g1)
>>>         return grad_input

save_for_backward(*tensors)

保存给定的张量以供将来调用 backward()。

save_for_backward 最多只能调用一次，并且只能在 forward() 方法内部调用，并且只能使用张量。

所有打算在反向传播中使用的张量都应该使用save_for_backward保存（而不是直接保存在ctx上），以防止梯度计算错误和内存泄漏，并启用保存张量钩子的应用。请参阅torch.autograd.graph.saved_tensors_hooks。

请注意，如果保存了中间张量（既不是输入也不是 forward() 的输出的张量）用于反向传播，您的自定义函数可能不支持双重反向传播。 不支持双重反向传播的自定义函数应使用 @once_differentiable 装饰其 backward() 方法，以便执行双重反向传播时引发错误。如果您希望支持双重反向传播，您可以根据反向传播期间的输入重新计算中间变量，或者将中间变量作为自定义函数的输出来返回。有关更多详细信息，请参阅 双重反向传播教程。

在 backward() 中，可以通过 saved_tensors 属性访问保存的张量。在将它们返回给用户之前，会进行检查以确保它们没有在任何就地操作中被修改。

参数也可以是 None。这是一个空操作。

有关如何使用此方法的更多详细信息，请参阅扩展 torch.autograd。

Example::

>>> 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 不是一个张量
>>>         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)

save_for_forward(*tensors)

保存给定的张量以供将来调用 jvp()。

save_for_forward 应该只在 forward() 方法内部调用一次，并且只能与张量一起调用。

在 jvp() 中，可以通过 saved_tensors 属性访问保存的对象。

参数也可以是 None。这是一个空操作。

有关如何使用此方法的更多详细信息，请参阅扩展 torch.autograd。

Example::

>>> 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)

set_materialize_grads(value)

设置是否具体化梯度张量。默认是 True。

这应该仅在 forward() 方法内部调用

如果 True，在调用 backward() 和 jvp() 方法之前，未定义的梯度张量将被扩展为全零张量。

Example::

>>> class SimpleFunc(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         return x.clone(), x.clone()
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):
>>>         return g1 + g2  # 不需要检查None
>>>
>>> # 我们修改SimpleFunc以处理非具体化的梯度输出
>>> 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:  # 现在我们必须检查None
>>>             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)  # 导致g2未定义

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 了解更多详情。

Return type

任意

static vjp(ctx, *grad_outputs)

定义一个用于反向模式自动微分的操作的微分公式。

此函数应由所有子类重写。 （定义此函数等同于定义 vjp 函数。）

它必须接受一个上下文 ctx 作为第一个参数，后面跟着与 forward() 返回的输出一样多的参数（对于 forward 函数中的非张量输出，将传递 None），并且它应该返回与 forward() 的输入一样多的张量。每个参数是相对于给定输出的梯度，每个返回值应该是相对于相应输入的梯度。如果输入不是张量或是不需要梯度的张量，你可以直接为该输入传递 None 作为梯度。

上下文可以用于检索在前向传递过程中保存的张量。它还具有一个属性 ctx.needs_input_grad，作为布尔元组，表示每个输入是否需要梯度。例如， backward() 将在 forward() 的第一个输入需要相对于输出的梯度计算时，设置 ctx.needs_input_grad[0] = True。

Return type

任意

static vmap(info, in_dims, *args)

在此处定义此 autograd.Function 的行为 torch.vmap()。

对于支持torch.autograd.Function()的torch.vmap()，你必须重写此静态方法，或者将generate_vmap_rule设置为True（你不能同时执行这两者）。

如果你选择覆盖这个静态方法：它必须接受

• 一个 info 对象作为第一个参数。info.batch_size 指定了被 vmapped 的维度的尺寸， 而 info.randomness 是传递给 torch.vmap() 的随机性选项。

• 作为第二个参数的 in_dims 元组。 对于 args 中的每个参数，in_dims 都有一个对应的 Optional[int]。如果参数不是张量或参数不被 vmapped 覆盖，则为 None，否则，它是一个整数 指定张量的哪个维度被 vmapped 覆盖。

• *args，这与传递给 forward() 的参数相同。

vmap静态方法的返回值是一个包含(output, out_dims)的元组。 类似于in_dims，out_dims应该与 output具有相同的结构，并且为每个输出包含一个out_dim，用于指定输出是否具有vmapped维度以及其索引位置。

请参阅使用 autograd.Function 扩展 torch.func了解更多详情。