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
它必须接受上下文 ctx 作为第一个参数,然后是任意数量的参数(张量或其他类型)。
有关更多详细信息,请参阅 组合或分开的 forward() 和 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 上存储张量(尽管出于向后兼容性考虑,目前没有强制执行此操作)。相反,应使用
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 的前向传递。
或者
用签名
forward(ctx, *args, **kwargs)
重写 forward。不重写setup_context
。在forward
内部设置用于反向传播的 ctx。用签名
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_dims
,out_dims
应该与output
具有相同的结构,并为每个输出包含一个out_dim
,用于指定输出是否具有 vmap 维度以及该维度在输出中的索引。更多详细信息,请参阅 使用 autograd.Function 扩展 torch.func。