InplaceFunction

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

该类仅出于向后兼容原因而存在。对于任何新的用例，请使用 Function 而不是这个。

static backward(ctx, *grad_outputs)

定义使用反向模式自动微分来区分操作的公式。

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

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

上下文可用于检索在 forward 传播期间保存的 tensor。它还有一个属性 ctx.needs_input_grad，它是一个布尔元组，表示每个输入是否需要梯度。例如，如果 forward() 的第一个输入需要计算相对于输出的梯度，那么 backward() 将有 ctx.needs_input_grad[0] = True。

返回类型

任何

static forward(*args, **kwargs)

定义自定义自动微分函数的前向传播。

此函数应被所有子类覆盖。定义 forward 有两种方式

用法 1 (组合 forward 和 ctx)

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

• 它必须接受一个 ctx 作为第一个参数，后跟任意数量的参数（tensor 或其他类型）。

• 有关更多详细信息，请参阅 组合或分离 forward() 和 setup_context()。

用法 2 (分离 forward 和 ctx)

@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 是 forward 的输出，inputs 是一个包含 forward 输入的元组。

• 有关更多详细信息，请参阅 扩展 torch.autograd。

上下文可用于存储可以在 backward 传播期间检索的任意数据。不应直接将 tensor 存储在 ctx 上（尽管出于向后兼容性原因，目前不强制执行）。相反，tensor 应使用 ctx.save_for_backward() 保存（如果它们打算在 backward 中使用（等价于 vjp））或使用 ctx.save_for_forward() 保存（如果它们打算在 jvp 中使用）。

返回类型

任何

static jvp(ctx, *grad_inputs)

定义使用前向模式自动微分来区分操作的公式。

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

You can use the ctx object to pass any value from the forward to this functions.

返回类型

任何

mark_dirty(*args)

将给定张量标记为在就地操作中已修改。

应在 setup_context() 或 forward() 方法中最多调用一次，所有参数都应该是输入。

在 forward() 调用中以原地方式修改的每个 tensor 都应传递给此函数，以确保我们的检查的正确性。函数是在修改之前还是之后调用的并不重要。

示例：

>>> class Inplace(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         x_npy = x.numpy() # x_npy shares storage with 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)  # This would lead to wrong gradients!
>>>                   # but the engine would not know unless we mark_dirty
>>> b.backward() # RuntimeError: one of the variables needed for gradient
>>>              # computation has been modified by an inplace operation

mark_non_differentiable(*args)

将输出标记为不可微分。

应在 setup_context() 或 forward() 方法中最多调用一次，所有参数都应该是 tensor 输出。

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

此功能用于例如从排序返回的索引。请参阅示例：

>>> 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):  # still need to accept 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 应在 setup_context() 或 forward() 方法中最多调用一次，并且只能使用 tensor。

所有打算在 backward 传播中使用但不是 forward 函数的输入或输出的 tensor 都应使用 save_for_backward 保存（而不是直接保存在 ctx 上），以防止梯度不正确和内存泄漏，并启用已保存 tensor hook 的应用。请参阅 torch.autograd.graph.saved_tensors_hooks。有关更多详细信息，请参阅 扩展 torch.autograd。

注意，如果保存了中间 tensor（不是 forward 函数的输入或输出的 tensor）用于 backward，您的自定义 Function 可能不支持 double backward。不支持 double backward 的自定义 Function 应该使用 @once_differentiable 装饰其 backward() 方法，以便执行 double backward 时会引发错误。如果您想支持 double backward，可以：在 backward 时根据输入重新计算中间项，或将中间项作为自定义 Function 的输出返回。有关更多详细信息，请参阅 double backward 教程。

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

参数也可以是 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)

save_for_forward(*tensors)

Save given tensors for a future call to jvp().

save_for_forward 应在 setup_context() 或 forward() 方法中最多调用一次，所有参数都应该是 tensor。

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

set_materialize_grads(value)

Set whether to materialize grad tensors. Default is True.

这应该只从 setup_context() 或 forward() 方法中调用。

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

示例

>>> 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 传播有两种方式。

Either

1. 使用签名 forward(ctx, *args, **kwargs) 覆盖 forward。不覆盖 setup_context。ctx 的设置用于 backward 在 forward 中完成。

2. 使用签名 forward(*args, **kwargs) 覆盖 forward 并覆盖 setup_context。ctx 的设置用于 backward 在 setup_context 中完成（而不是在 forward 中）。

有关更多详细信息，请参阅 torch.autograd.Function.forward() 和 扩展 torch.autograd。

返回类型

任何

static vjp(ctx, *grad_outputs)

定义使用反向模式自动微分来区分操作的公式。

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

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

上下文可用于检索在 forward 传播期间保存的 tensor。它还有一个属性 ctx.needs_input_grad，它是一个布尔元组，表示每个输入是否需要梯度。例如，如果 forward() 的第一个输入需要计算相对于输出的梯度，那么 backward() 将有 ctx.needs_input_grad[0] = True。

返回类型

任何

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 中的每个 arg，in_dims 有一个相应的 Optional[int]。如果 arg 不是 Tensor 或 arg 不被 vmap，则为 None，否则，它是一个指定 Tensor 的哪个维度被 vmap 的整数。

• *args，与 forward() 的 args 相同。

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

有关更多详细信息，请参阅 使用 autograd.Function 扩展 torch.func。