InplaceFunction

类 torch.autograd.function.InplaceFunction(inplace=False)[source][source] ¶

此类仅用于向后兼容。对于任何新的用例，请使用 Function 替代此功能。

static backward(ctx, *grad_outputs)

定义使用反向模式自动微分对操作进行微分公式的公式。

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

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

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

返回类型:

任何

static forward(*args, **kwargs)

定义自定义 autograd Function 的前向操作。

此函数需由所有子类重写。定义前向传播有两种方式：

使用方法 1（组合前向和上下文）：

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

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

• 更多详情请参阅组合或分离的前向()和 setup_context()。

使用说明 2（分离前向和 ctx）：

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

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

• 前向操作不再接受 ctx 参数。

• 取而代之，您还必须重写 torch.autograd.Function.setup_context() 静态方法来处理设置 ctx 对象。 output 是前向操作的输出， inputs 是前向操作的输入元组。

• 更多详情请参阅扩展 torch.autograd。

上下文可以用来存储任意数据，这些数据可以在反向传播过程中被检索。不应直接在 ctx 上存储张量（尽管为了向后兼容，目前没有强制执行）。相反，如果张量打算用于 backward （等价于 vjp ）或 ctx.save_for_forward() 如果打算用于 jvp ，则应使用 ctx.save_for_backward() 保存。

返回类型:

任何

static jvp(ctx, *grad_inputs)

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

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

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

返回类型:

任何

mark_dirty(*args)

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

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

在调用 forward() 中对原地修改过的每个张量都应传递给此函数，以确保检查的正确性。函数是在修改前后调用无关紧要。

示例::

>>> 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() 方法中，所有参数都应该是张量输出。

这将标记输出为不需要梯度，从而提高反向计算的效率。你仍然需要在 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() 。

应最多调用一次，在 setup_context() 或 forward() 方法中，并且仅与张量一起调用。

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

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

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

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)

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)[source]

定义 autograd.Function 的前向传递有两种方式。

要么：

1. 使用签名 forward(ctx, *args, **kwargs) 覆盖前进。 setup_context 不会被覆盖。设置反向的 ctx 在 forward 内部进行。

2. 使用签名 forward(*args, **kwargs) 和覆盖 setup_context 。设置反向的 ctx 在 setup_context （而不是在 forward 内部）进行。

请参阅 torch.autograd.Function.forward() 和扩展 torch.autograd 以获取更多详细信息。

返回类型:

任何

static vjp(ctx, *grad_outputs)

定义使用反向模式自动微分对操作进行微分公式的公式。

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

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

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

返回类型:

任何

static vmap(info, in_dims, *args)

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

为了支持 torch.autograd.Function() ，您必须重写此静态方法，或将 generate_vmap_rule 设置为 True （您不能同时进行这两项操作）。

如果您选择重写此静态方法：它必须接受

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

• 使用 in_dims 元组作为第二个参数。对于 args 中的每个参数， in_dims 都有一个对应的 Optional[int] 。如果参数不是 Tensor 或者参数没有被 vmapped，则是 None ，否则，它是一个整数，指定了 Tensor 被 vmapped 的维度。

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

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

请参阅《使用 autograd.Function 扩展 torch.func 的更多详细信息。》