快捷方式

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,它是一个布尔值的元组,表示每个输入是否需要梯度。例如,如果 forward() 的第一个输入需要计算关于输出的梯度,则 backward() 将具有 ctx.needs_input_grad[0] = True

返回类型

任何

static forward(*args, **kwargs)

定义自定义 autograd Function 的正向传播。

所有子类都必须重写此函数。有两种方法可以定义正向传播

用法 1(组合正向传播和 ctx)

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

用法 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 上存储张量(尽管出于向后兼容性考虑,目前没有强制执行此操作)。相反,应使用 ctx.save_for_backward() 保存张量,如果打算在 backward(等效地,vjp)中使用它们,或者使用 ctx.save_for_forward() 保存张量,如果打算在 jvp 中使用它们。

返回类型

任何

static jvp(ctx, *grad_inputs)

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

此函数需要由所有子类重写。它必须接受一个上下文 ctx 作为第一个参数,后面跟着与 forward() 获取的输入数量相同的输入(对于 forward 函数的非张量输入将传递 None),并且它应该返回与 forward() 的输出数量相同的张量。每个参数是相对于给定输入的梯度,每个返回值应该是相对于相应输出的梯度。如果输出不是张量或函数相对于该输出不可微,则可以将 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()

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

有两种方法可以定义 autograd.Function 的前向传递。

或者

  1. 用签名 forward(ctx, *args, **kwargs) 重写 forward。不重写 setup_context。在 forward 内部设置用于反向传播的 ctx。

  2. 用签名 forward(*args, **kwargs) 重写 forward 并重写 setup_context。在 setup_context 内部设置用于反向传播的 ctx(而不是在 forward 内部设置)。

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

返回类型

任何

static vjp(ctx, *grad_outputs)

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

所有子类都必须重写此函数。(定义此函数等效于定义 vjp 函数。)

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

上下文可用于检索在正向传递期间保存的张量。它还有一个属性 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 中的每个参数,in_dims 都有一个对应的 Optional[int]。如果参数不是张量或没有进行 vmap 操作,则为 None;否则,它是一个整数,指定了张量的哪个维度正在进行 vmap 操作。

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

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

更多详细信息,请参阅 使用 autograd.Function 扩展 torch.func

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源