InplaceFunction

类torch.autograd.function.InplaceFunction(inplace=False)[源代码]

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

static backward(ctx, *grad_outputs)

定义一个公式，用于通过反向模式自动微分来区分操作。

此函数需要在所有子类中被重写。（定义此函数相当于定义了vjp函数。）

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

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

返回类型

Any

静态前传(*args, **kwargs)

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

此函数需要被所有子类重写。有兩種方法可以定義前向傳播：

用法 1（结合前向和 ctx）：

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

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

• 详情请参见结合或分开使用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

• “前向”函数不再接收ctx参数。

• 相反，你还必须重写 torch.autograd.Function.setup_context() 静态方法来设置 ctx 对象。其中，output 是前向传播的输出，而 inputs 是一个包含前向传播输入的元组。

• 更多细节请参见扩展torch.autograd

上下文可以用来存储任意数据，并在反向传递期间检索这些数据。张量不应直接存储在ctx上（尽管当前还没有强制执行以保持后向兼容性）。相反，如果张量用于backward（等同于vjp），则应使用ctx.save_for_backward()进行保存；如果张量用于jvp，则应使用ctx.save_for_forward()进行保存。

返回类型

Any

staticjvp(ctx, *grad_inputs)

定义一个公式，用于使用前向模式自动微分来计算操作的导数。

此函数需要被所有子类重写。它必须接受一个上下文 ctx 作为第一个参数，然后是与 forward() 方法接收到的输入数量相同的参数（对于前向函数中非张量类型的输入将传递 None）。该方法应返回与 forward() 输出相同数量的张量。每个参数是相对于给定输入的梯度，每个返回值应该是相对于相应输出的梯度。如果某个输出不是 Tensor 或者函数对该输出不可微，则可以将 None 作为该输入的梯度传递。

你可以使用ctx对象将任何值从前进方向传递到此函数。

返回类型

Any

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()保存给定的张量。

save_for_backward 应该最多调用一次，在 setup_context() 或者 forward() 方法中，并且只能传递张量。

所有计划在反向传递中使用的张量应使用save_for_backward（而不是直接保存到ctx）来保存，以防止计算不正确的梯度和内存泄漏，并启用已保存张量挂钩的应用。参见torch.autograd.graph.saved_tensors_hooks。

注意，如果在反向传播中保存了既不是输入也不是输出的中间张量，你的自定义函数可能不支持双次反向传播。对于不支持双次反向传播的自定义函数，应该用@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(值)

设置是否将梯度张量进行具体化。默认值为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。在 forward 方法中为反向传播设置 ctx。

2. 重写带有签名 forward(*args, **kwargs) 的 forward 方法，并且重写 setup_context 方法。设置 ctx 以供反向传播使用应在 setup_context 方法中进行（而不是在 forward 方法中）。

了解更多详情，请参阅torch.autograd.Function.forward() 和 扩展 torch.autograd。

返回类型

Any

staticvjp(ctx, *grad_outputs)

定义一个公式，用于通过反向模式自动微分来区分操作。

此函数需要在所有子类中被重写。（定义此函数相当于定义了vjp函数。）

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

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

返回类型

Any

staticvmap(info, in_dims, *args)

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

为了使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，以指定该输出是否具有 vmap 维度及其索引。

请参见使用 autograd.Function 扩展 torch.func以获取更多详细信息。