扩展 PyTorch ¶
在这篇笔记中，我们将介绍扩展 torch.nn ， torch.autograd ， torch 以及编写自定义 C++ 扩展的方法。
添加新的算子 ¶
PyTorch 提供了一个大型算子库，这些算子可以在张量上工作（例如 torch.add() ， torch.sum() 等）。然而，您可能希望将一个新的自定义操作引入 PyTorch，并使其表现得像 PyTorch 的内置算子一样。为此，您必须通过 Python 的 torch.library 或 C++的 TORCH_LIBRARY API 将自定义操作注册到 PyTorch 中。
请参阅 PyTorch 自定义算子介绍页面以获取更多详细信息。
扩展 torch.autograd ¶
向 autograd 添加操作需要为每个操作实现一个新的 Function 子类。回想一下，函数是 autograd 使用来编码操作历史和计算梯度的。
这份文档的第一部分专注于反向模式自动微分，因为它是使用最广泛的功能。文档末尾讨论了正向模式自动微分的扩展。
何时使用 ¶
通常情况下，如果您想在模型中执行不可微分的计算或依赖于非 PyTorch 库（例如 NumPy）的操作，但仍然希望您的操作可以与其他操作链式调用并使用 autograd 引擎，则应实现自定义函数。
在某些情况下，自定义函数也可以用来提高性能和内存使用：如果您使用 C++扩展实现了正向和反向传递，则可以将它们包装在 Function 中以与 autograd 引擎接口。如果您想减少反向传递保存的缓冲区数量，可以使用自定义函数将操作组合在一起。
何时不使用 ¶
如果您已经可以用 PyTorch 的内置操作来编写您的函数，那么其反向图（很可能）已经能够被 autograd 记录。在这种情况下，您不需要自己实现反向函数。考虑使用普通的 Python 函数。
如果您需要维护状态，即训练参数，您也应该（同样）使用自定义模块。请参阅下文有关扩展 torch.nn 的更多信息。
如果您想在反向传播过程中修改梯度或执行副作用，请考虑注册一个张量或 Module 钩子。
如何使用 ¶
执行以下步骤：1. 继承 Function 并实现 forward() 方法，(可选) setup_context() 和 backward() 方法。2. 在 ctx 参数上调用适当的方法。3. 声明你的函数是否支持双向回溯。4. 使用 gradcheck 验证你的梯度是否正确。
步骤 1：在继承 Function 之后，你需要定义 3 个方法：
forward()是执行操作的代码。它可以接受任意数量的参数，其中一些参数是可选的，如果你指定了默认值。这里接受所有类型的 Python 对象。Tensor参数跟踪历史（即，带有requires_grad=True的参数）将在调用之前转换为不跟踪历史的参数，并且它们的使用将记录在图中。请注意，此逻辑不会遍历列表/字典/任何其他数据结构，而只会考虑调用中的直接参数张量。你可以返回单个Tensor输出，或者如果有多个输出，则返回一个tuple张量数组。另外，请参阅Function的文档，以找到只能从forward()调用的有用方法的描述。setup_context()（可选）。可以编写一个接受forward()对象的“组合”ctx，或者（从 PyTorch 2.0 开始）一个不接收forward()的单独ctx，以及一个setup_context()方法，其中ctx的修改发生。forward()应该负责计算，而setup_context()只应负责ctx的修改（而不进行任何计算）。一般来说，单独的forward()和setup_context()更接近 PyTorch 原生操作的方式，因此与各种 PyTorch 子系统更易于组合。有关详细信息，请参阅“组合或分离 forward()和 setup_context()”。backward()（或vjp()）定义了梯度公式。它将根据输出数量提供Tensor个参数，每个参数代表相对于该输出的梯度。重要的是永远不要就地修改这些。它应该返回与输入数量相同的张量，每个张量包含相对于其对应输入的梯度。如果您的输入不需要梯度（needs_input_grad是一个布尔值元组，指示每个输入是否需要梯度计算），或者是非Tensor对象，则可以返回python:None。如果forward()有可选参数，则可以返回比输入更多的梯度，只要它们都是None。
步骤 2：确保新功能 Function 正确使用 ctx 的函数，以使新功能与自动微分引擎正常工作。
save_for_backward()必须用于保存任何需要在反向传播中使用的张量。非张量应直接存储在 ctx 上。如果保存了既不是输入也不是输出的张量进行反向传播，则你的Function可能不支持双重反向（见步骤 3）。mark_dirty()必须用于标记任何被前向函数就地修改的输入。mark_non_differentiable()必须用于告知引擎某个输出是否不可微分。默认情况下，所有可微分的输出张量都将设置为需要梯度。非可微分类型（即整型）的张量永远不会标记为需要梯度。set_materialize_grads()可以用来告诉自动微分引擎在输出不依赖于输入的情况下优化梯度计算，即不将给 backward 函数的 grad 张量实体化。也就是说，如果设置为 False，Python 中的 None 对象或 C++ 中的“未定义张量”（x.defined() 为 False 的张量 x）在调用 backward 之前不会转换为填充零的张量，因此您的代码需要将这些对象作为填充零的张量来处理。此设置的默认值为 True。
步骤 3：如果您的 Function 不支持双反向，则应通过装饰 backward 使用 once_differentiable() 显式声明。使用此装饰器，尝试通过您的函数执行双反向将产生错误。有关双反向的更多信息，请参阅我们的双反向教程。
步骤 4：建议您使用 torch.autograd.gradcheck() 检查您的 backward 函数是否正确计算了前向的梯度，通过使用您的 backward 函数计算雅可比矩阵，并逐元素与使用有限差分法数值计算的雅可比矩阵进行比较。
示例
下面您可以找到 Linear 函数的代码，其中包含额外的注释：
# Inherit from Function
class LinearFunction(Function):
# Note that forward, setup_context, and backward are @staticmethods
@staticmethod
def forward(input, weight, bias):
output = input.mm(weight.t())
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
return output
@staticmethod
# inputs is a Tuple of all of the inputs passed to forward.
# output is the output of the forward().
def setup_context(ctx, inputs, output):
input, weight, bias = inputs
ctx.save_for_backward(input, weight, bias)
# This function has only a single output, so it gets only one gradient
@staticmethod
def backward(ctx, grad_output):
# This is a pattern that is very convenient - at the top of backward
# unpack saved_tensors and initialize all gradients w.r.t. inputs to
# None. Thanks to the fact that additional trailing Nones are
# ignored, the return statement is simple even when the function has
# optional inputs.
input, weight, bias = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None
# These needs_input_grad checks are optional and there only to
# improve efficiency. If you want to make your code simpler, you can
# skip them. Returning gradients for inputs that don't require it is
# not an error.
if ctx.needs_input_grad[0]:
grad_input = grad_output.mm(weight)
if ctx.needs_input_grad[1]:
grad_weight = grad_output.t().mm(input)
if bias is not None and ctx.needs_input_grad[2]:
grad_bias = grad_output.sum(0)
return grad_input, grad_weight, grad_bias
现在为了让这些自定义操作更容易使用，我们建议要么给它们起别名，要么将它们封装在函数中。将它们封装在函数中可以让我们支持默认参数和关键字参数：
# Option 1: alias
linear = LinearFunction.apply
# Option 2: wrap in a function, to support default args and keyword args.
def linear(input, weight, bias=None):
return LinearFunction.apply(input, weight, bias)
这里，我们给出一个由非 Tensor 参数参数化的函数的额外示例：
class MulConstant(Function):
@staticmethod
def forward(tensor, constant):
return tensor * constant
@staticmethod
def setup_context(ctx, inputs, output):
# ctx is a context object that can be used to stash information
# for backward computation
tensor, constant = inputs
ctx.constant = constant
@staticmethod
def backward(ctx, grad_output):
# We return as many input gradients as there were arguments.
# Gradients of non-Tensor arguments to forward must be None.
return grad_output * ctx.constant, None
这里，我们通过调用 set_materialize_grads(False)优化了上面的示例：
class MulConstant(Function):
@staticmethod
def forward(tensor, constant):
return tensor * constant
@staticmethod
def setup_context(ctx, inputs, output):
tensor, constant = inputs
ctx.set_materialize_grads(False)
ctx.constant = constant
@staticmethod
def backward(ctx, grad_output):
# Here we must handle None grad_output tensor. In this case we
# can skip unnecessary computations and just return None.
if grad_output is None:
return None, None
# We return as many input gradients as there were arguments.
# Gradients of non-Tensor arguments to forward must be None.
return grad_output * ctx.constant, None
如果需要在 forward() 中计算的任何“中间”Tensor 被保存，它们必须作为输出返回，或者将 forward 和 setup_context() （参见 Combined 或 separate forward()和 setup_context()）组合起来。请注意，这意味着如果您想让梯度通过这些中间值流动，您需要为它们定义梯度公式（也请参阅双反向教程）：
class MyCube(torch.autograd.Function):
@staticmethod
def forward(x):
# We wish to save dx for backward. In order to do so, it must
# be returned as an output.
dx = 3 * x ** 2
result = x ** 3
return result, dx
@staticmethod
def setup_context(ctx, inputs, output):
x, = inputs
result, dx = output
ctx.save_for_backward(x, dx)
@staticmethod
def backward(ctx, grad_output, grad_dx):
x, dx = ctx.saved_tensors
# In order for the autograd.Function to work with higher-order
# gradients, we must add the gradient contribution of `dx`,
# which is grad_dx * 6 * x.
result = grad_output * dx + grad_dx * 6 * x
return result
# Wrap MyCube in a function so that it is clearer what the output is
def my_cube(x):
result, dx = MyCube.apply(x)
return result
注意
输入到 backward ，即 grad_output ，也可以是跟踪历史的张量。所以如果 backward 使用可微操作实现（例如调用另一个自定义的 Function ），高阶导数将正常工作。在这种情况下，使用 save_for_backward 保存的张量也可以用于反向传播，并且梯度可以反向流动，但保存到 ctx 的张量则不会有梯度反向流动。如果您需要使保存到 ctx 的张量有梯度反向流动，您应该将其作为自定义 Function 的输出并使用 save_for_backward 保存。
您可能想检查您实现的反向方法是否实际上计算了您函数的导数。这可以通过使用小有限差分与数值近似进行比较来实现：
from torch.autograd import gradcheck
# gradcheck takes a tuple of tensors as input, check if your gradient
# evaluated with these tensors are close enough to numerical
# approximations and returns True if they all verify this condition.
input = (torch.randn(20,20,dtype=torch.double,requires_grad=True), torch.randn(30,20,dtype=torch.double,requires_grad=True))
test = gradcheck(linear, input, eps=1e-6, atol=1e-4)
print(test)
更多关于有限差分梯度比较的细节请参阅数值梯度检查。如果您的函数用于高阶导数（对反向传播进行微分）中，您可以使用同一包中的 gradgradcheck 函数来检查高阶导数。
组合或分离的 forward() 和 setup_context() ¶
定义 Function 主要有两种方式。要么：
将前向计算逻辑与
setup_context()结合定义一个forward()（截至 PyTorch 2.0）分别定义一个
forward()和setup_context()
我们推荐第二种方案（分别定义 forward() 和 setup_context() ），因为这更接近 PyTorch 原生操作的实现方式，并且可以与 torch.func 转换器进行组合。然而，我们计划在将来支持两种方法；将 forward() 与 setup_context() 结合：由于可以保存中间结果而不必作为输出返回，因此这提供了更多的灵活性。
请参阅上一节了解如何定义 Function ，其中包含独立的 forward() 和 setup_context() 。
下面是一个如何定义包含合并的 forward() 和 setup_context() 的 Function 的例子：
class LinearFunction(Function):
@staticmethod
# ctx is the first argument to forward
def forward(ctx, input, weight, bias=None):
# The forward pass can use ctx.
ctx.save_for_backward(input, weight, bias)
output = input.mm(weight.t())
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
return output
@staticmethod
def backward(ctx, grad_output):
input, weight, bias = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None
if ctx.needs_input_grad[0]:
grad_input = grad_output.mm(weight)
if ctx.needs_input_grad[1]:
grad_weight = grad_output.t().mm(input)
if bias is not None and ctx.needs_input_grad[2]:
grad_bias = grad_output.sum(0)
return grad_input, grad_weight, grad_bias
前向模式 AD ¶
覆盖前向模式 AD 公式具有非常相似的 API，但也有一些细微差别。您可以实现 jvp() 函数。
将提供与输入数量相等的 Tensor 参数，每个参数代表相对于该输入的梯度。应返回与输出数量相等的张量，每个张量包含相对于其对应输出的梯度。 jvp() 将在 forward() 方法之后、 apply() 返回之前被调用。
jvp() 与 backward() 函数有一些细微的区别：
可以使用 ctx 从
forward()传递任何数据到jvp()函数。如果该状态对于backward()不需要，可以在jvp()函数结束时显式释放它，执行del ctx.foo。jvp()的实现必须是可反向求导的，或者显式检查给定的正向模式梯度是否没有设置requires_grad。函数
jvp()必须匹配视图/inplace 行为forward()。例如，如果第i个输入被就地修改，那么第i个梯度必须就地更新。同样地，如果第j个输出是第k个输入的视图。那么返回的第j个输出梯度必须是给定第k个输入梯度的视图。因为用户无法指定需要计算哪个梯度，所以
jvp()函数应该始终计算所有输出的梯度。前向模式梯度会尊重
set_materialize_grads()设置的标志，并且当此功能禁用时，您可以获取 None 输入梯度。
torch.func 转换和/或 torch.vmap() ¶
请参阅使用 autograd.Function 扩展 torch.func 的详细信息。
torch.nn ¶
nn 导出两种接口 - 模块及其功能版本。您可以通过这两种方式扩展它，但我们建议对于所有持有任何参数或缓冲区的层，使用模块，并建议使用无参数的功能形式进行操作，如激活函数、池化等。
操作的功能版本已在上述章节中完全介绍。
添加一个 Module ¶
由于 nn 严重依赖 autograd ，添加一个新 Module 需要实现一个执行操作并能计算梯度的 Function 。从现在开始，让我们假设我们想要实现一个 Linear 模块，并且我们已经实现了如上所示的功能。添加这个功能所需的代码非常少。现在，需要实现两个函数：
__init__（可选）- 接受诸如内核大小、特征数量等参数并初始化参数和缓冲区。forward()- 实例化一个Function并使用它来执行操作。它与上面显示的功能包装器非常相似。
这就是如何实现一个 Linear 模块的示例：
class Linear(nn.Module):
def __init__(self, input_features, output_features, bias=True):
super().__init__()
self.input_features = input_features
self.output_features = output_features
# nn.Parameter is a special kind of Tensor, that will get
# automatically registered as Module's parameter once it's assigned
# as an attribute. Parameters and buffers need to be registered, or
# they won't appear in .parameters() (doesn't apply to buffers), and
# won't be converted when e.g. .cuda() is called. You can use
# .register_buffer() to register buffers.
# nn.Parameters require gradients by default.
self.weight = nn.Parameter(torch.empty(output_features, input_features))
if bias:
self.bias = nn.Parameter(torch.empty(output_features))
else:
# You should always register all possible parameters, but the
# optional ones can be None if you want.
self.register_parameter('bias', None)
# Not a very smart way to initialize weights
nn.init.uniform_(self.weight, -0.1, 0.1)
if self.bias is not None:
nn.init.uniform_(self.bias, -0.1, 0.1)
def forward(self, input):
# See the autograd section for explanation of what happens here.
return LinearFunction.apply(input, self.weight, self.bias)
def extra_repr(self):
# (Optional)Set the extra information about this module. You can test
# it by printing an object of this class.
return 'input_features={}, output_features={}, bias={}'.format(
self.input_features, self.output_features, self.bias is not None
)
扩展 torch Python API ¶
您可以通过定义一个具有匹配 Tensor 方法的自定义类来创建模拟 Tensor 的自定义类型。但您想将这些类型传递给像 torch.add() 这样的函数，这些函数位于顶层 torch 命名空间中，并接受 Tensor 操作数呢？
如果您的自定义 Python 类型定义了一个名为 __torch_function__ 的方法，当您的自定义类的实例传递给 torch 命名空间中的函数时，PyTorch 将调用您的 __torch_function__ 实现。这使得您可以定义任何 torch 命名空间中函数的定制实现，您的 __torch_function__ 实现可以调用这些函数，允许您的用户使用您自定义的类型与现有 PyTorch 工作流程相结合，这些工作流程是他们已经为 Tensor 编写的。这也适用于与 Tensor 无关的“鸭子”类型以及用户定义的 Tensor 子类。
使用类似 Tensor 的类型扩展 torch -
注意
此功能灵感来源于 NumPy 的 __array_function__ 协议。有关更多详细信息，请参阅 NumPy 文档和 NEP-0018。
为了具体说明，让我们从一个简单的例子开始，该例子说明了 API 分派机制。我们将创建一个自定义类型，该类型表示一个二维标量张量，由顺序 N 和对角线元素的值 value 参数化：
class ScalarTensor(object):
def __init__(self, N, value):
self._N = N
self._value = value
def __repr__(self):
return "ScalarTensor(N={}, value={})".format(self._N, self._value)
def tensor(self):
return self._value * torch.eye(self._N)
此设计的第一个迭代版本并不十分有用。 ScalarTensor 的主要功能是提供比基类张量更紧凑的标量张量字符串表示形式：
>>> d = ScalarTensor(5, 2)
>>> d
ScalarTensor(N=5, value=2)
>>> d.tensor()
tensor([[2., 0., 0., 0., 0.],
[0., 2., 0., 0., 0.],
[0., 0., 2., 0., 0.],
[0., 0., 0., 2., 0.],
[0., 0., 0., 0., 2.]])
如果我们尝试使用此对象与 torch API 一起使用，我们将遇到问题：
>>> import torch
>>> torch.mean(d)
TypeError: mean(): argument 'input' (position 1) must be Tensor, not ScalarTensor
将 __torch_function__ 实现添加到 ScalarTensor 中可以使上述操作成功。让我们重新实现我们的实现，这次添加一个 __torch_function__ 实现：
HANDLED_FUNCTIONS = {}
class ScalarTensor(object):
def __init__(self, N, value):
self._N = N
self._value = value
def __repr__(self):
return "ScalarTensor(N={}, value={})".format(self._N, self._value)
def tensor(self):
return self._value * torch.eye(self._N)
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
if func not in HANDLED_FUNCTIONS or not all(
issubclass(t, (torch.Tensor, ScalarTensor))
for t in types
):
return NotImplemented
return HANDLED_FUNCTIONS[func](*args, **kwargs)
__torch_function__ 方法接受四个参数： func ，被覆盖的 torch API 函数的引用， types ，实现 __torch_function__ 的 Tensor-like 类型的列表， args ，传递给函数的参数元组，以及 kwargs ，传递给函数的关键字参数字典。它使用名为 HANDLED_FUNCTIONS 的全局调度表来存储自定义实现。该字典的键是 torch 命名空间中的函数，值是 ScalarTensor 的实现。
注意
使用全局调度表不是 __torch_function__ API 的要求，它只是结构化你的覆盖实现的有用设计模式。
这个类定义还不够，使得当我们传递一个 ScalarTensor 给 torch.mean 时，它不能正确地执行——我们还需要为 ScalarTensor 操作数定义一个 torch.mean 的实现，并将其添加到 HANDLED_FUNCTIONS 分派表字典中。实现这一目标的一种方法是为我们的覆盖实现定义一个装饰器：
import functools
def implements(torch_function):
"""Register a torch function override for ScalarTensor"""
def decorator(func):
functools.update_wrapper(func, torch_function)
HANDLED_FUNCTIONS[torch_function] = func
return func
return decorator
这可以应用于我们的覆盖实现的实现：
@implements(torch.mean)
def mean(input):
return float(input._value) / input._N
通过这个更改，我们现在可以使用 torch.mean 和 ScalarTensor ：
>>> d = ScalarTensor(5, 2)
>>> torch.mean(d)
0.4
当然， torch.mean 是一种最简单的覆盖函数的例子，因为它只接受一个操作数。我们可以使用相同的机制来覆盖接受一个以上操作数的函数，其中任何一个操作数可能是一个定义了 __torch_function__ 的张量或类似张量，例如用于 torch.add() ：
def ensure_tensor(data):
if isinstance(data, ScalarTensor):
return data.tensor()
return torch.as_tensor(data)
@implements(torch.add)
def add(input, other):
try:
if input._N == other._N:
return ScalarTensor(input._N, input._value + other._value)
else:
raise ValueError("Shape mismatch!")
except AttributeError:
return torch.add(ensure_tensor(input), ensure_tensor(other))
本版本在两个操作数都是 ScalarTensor 实例时具有快速路径，同时还有一个较慢的路径，当任一操作数不是 ScalarTensor 时退化到将数据转换为张量。这使得当任一操作数是 ScalarTensor 或常规的 Tensor 时，覆盖函数能够正确执行：
>>> s = ScalarTensor(2, 2)
>>> torch.add(s, s)
ScalarTensor(N=2, value=4)
>>> t = torch.tensor([[1, 1,], [1, 1]])
>>> torch.add(s, t)
tensor([[3., 1.],
[1., 3.]])
注意，我们的 add 实现不像 torch.add() 那样将 alpha 或 out 作为关键字参数：
>>> torch.add(s, s, alpha=2)
TypeError: add() got an unexpected keyword argument 'alpha'
为了速度和灵活性， __torch_function__ 调度机制不会检查覆盖函数的签名是否与被覆盖函数在 torch API 中的签名匹配。对于某些应用程序，忽略可选参数是可以接受的，但为了确保与 Tensor 完全兼容，用户实现的 torch API 函数应确保精确模拟被覆盖函数的 API。
torch API 中未明确覆盖的函数将从 __torch_function__ 返回 NotImplemented 。如果所有带有 __torch_function__ 定义的操作数都返回 NotImplemented ，PyTorch 将引发 TypeError 。这意味着大多数情况下，没有为特定类型明确覆盖的操作将在传递此类类型的实例时引发 TypeError 。
>>> torch.mul(s, 3)
TypeError: no implementation found for 'torch.mul' on types that
implement __torch_function__: [ScalarTensor]
在实践中这意味着，如果您想使用 __torch_function__ 实现来实施覆盖，您需要明确实现完整的 torch API 或您用例中关心的整个 API 子集。这可能是一项艰巨的任务，因为完整的 torch API 相当广泛。
另一种选择是对于未处理的操作不返回 NotImplemented ，而是当没有覆盖时将 Tensor 传递给原始的 torch 函数。例如，如果我们将 __torch_function__ 的 ScalarTensor 实现更改为以下内容：
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
if func not in HANDLED_FUNCTIONS or not all(
issubclass(t, (torch.Tensor, ScalarTensor))
for t in types
):
args = [a.tensor() if hasattr(a, 'tensor') else a for a in args]
return func(*args, **kwargs)
return HANDLED_FUNCTIONS[func](*args, **kwargs)
那么 torch.mul() 将正常工作，尽管返回类型始终是 Tensor 而不是 ScalarTensor ，即使两个操作数都是 ScalarTensor 实例：
>>> s = ScalarTensor(2, 2)
>>> torch.mul(s, s)
tensor([[4., 0.],
[0., 4.]])
还请参阅下面的 MetadataTensor 示例，了解此模式的另一种变体，但始终返回 MetadataTensor 以通过 torch API 中的操作传播元数据。
__torch_function__ 协议旨在全面覆盖 API，部分覆盖可能导致不理想的结果，特别是某些函数会引发 TypeError 。这对于子类尤其如此，即使它们返回的结果完全相同，也必须覆盖 torch.add、torch.Tensor.__add__ 和 torch.Tensor.add 这三个。未能做到这一点可能会导致无限递归。如果需要从 torch.Tensor 子类中实现函数，必须在实现中使用 super().__torch_function__ 。
子类化 torch.Tensor
自 1.7.0 版本起， torch.Tensor 上的方法和应用于 torch.Tensor 子类的公共 torch.* 命名空间中的函数将返回子类实例而不是 torch.Tensor 实例：
>>> class SubTensor(torch.Tensor):
... pass
>>> type(torch.add(SubTensor([0]), SubTensor([1]))).__name__
'SubTensor'
>>> type(torch.add(SubTensor([0]), torch.tensor([1]))).__name__
'SubTensor'
如果存在多个子类，则默认选择层次结构中最低的子类。如果没有唯一的方法来确定这种情况，则引发 TypeError 异常：
>>> type(torch.add(SubTensor2([0]), SubTensor([1]))).__name__
'SubTensor2'
>>> type(torch.add(SubTensor2([0]), torch.tensor([1]))).__name__
'SubTensor2'
>>> torch.add(SubTensor([0]), OtherSubTensor([1]))
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: no implementation found for 'torch.add' on types that implement __torch_function__: [SubTensor, OtherSubTensor]
如果希望对所有张量方法进行全局覆盖，可以使用 __torch_function__ 。以下是一个记录所有函数/方法调用的示例：
class LoggingTensor(torch.Tensor):
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
# NOTE: Logging calls Tensor.__repr__, so we can't log __repr__ without infinite recursion
if func is not torch.Tensor.__repr__:
logging.info(f"func: {func.__name__}, args: {args!r}, kwargs: {kwargs!r}")
if kwargs is None:
kwargs = {}
return super().__torch_function__(func, types, args, kwargs)
然而，如果希望覆盖 Tensor 子类的方法，可以通过直接覆盖该方法（为子类定义它）或使用 __torch_function__ 并与 func 匹配来实现。
应当注意，在 __torch_function__ 中对于子类，始终调用 super().__torch_function__(func, ...) 而不是直接调用 func ，正如在版本 1.7.0 之前的情况。未能这样做可能会导致 func 递归回 __torch_function__ ，从而引起无限递归。
使用 Tensor 包装类型扩展 torch
另一个有用的例子是包装 Tensor 的类型，无论是作为属性还是通过子类化。下面我们实现这种类型的一个特殊情况，即 MetadataTensor ，它将元数据字典附加到通过 torch 操作传播的 Tensor 。由于这是一种通用的包装类型，用于完整的 torch API，因此我们不需要单独实现每个覆盖，这样可以使 __torch_function__ 的实现对允许的操作更加宽容：
class MetadataTensor(object):
def __init__(self, data, metadata=None, **kwargs):
self._t = torch.as_tensor(data, **kwargs)
self._metadata = metadata
def __repr__(self):
return "Metadata:\n{}\n\ndata:\n{}".format(self._metadata, self._t)
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
metadatas = tuple(a._metadata for a in args if hasattr(a, '_metadata'))
args = [getattr(a, '_t', a) for a in args]
assert len(metadatas) > 0
ret = func(*args, **kwargs)
return MetadataTensor(ret, metadata=metadatas[0])
这种简单的实现并不一定适用于 torch API 中的每个函数，但它足以捕获大多数常见操作：
>>> metadata = {'owner': 'Ministry of Silly Walks'}
>>> m = MetadataTensor([[1, 2], [3, 4]], metadata=metadata)
>>> t = torch.tensor([[1, 2], [1, 2]])
>>> torch.add(t, m)
Metadata:
{'owner': 'Ministry of Silly Walks'}
data:
tensor([[2, 4],
[4, 6]])
>>> torch.mul(t, m)
Metadata:
{'owner': 'Ministry of Silly Walks'}
data:
tensor([[1, 4],
[3, 8]])
对定义 __torch_function__ 的多个类型的操作
可以使用 torch API 与多个具有 __torch_function__ 实现的独立类型，但必须特别注意。在这种情况下，规则是：
调度操作收集每个操作数所有不同的
__torch_function__实现，并按顺序调用它们：先子类后父类，并在运算符表达式中从左到右。如果返回的不是
NotImplemented以外的任何值，则该值作为结果返回。实现可以注册它们不实现操作，通过返回NotImplemented来完成。如果所有的
__torch_function__实现都返回NotImplemented，PyTorch 将引发一个TypeError。
PyTorch API 覆盖测试覆盖率
实现 __torch_function__ 的一个麻烦之处在于，如果某些操作有覆盖而其他操作没有，用户最多只能看到不一致的体验，最坏的情况是在运行时使用没有覆盖的功能时引发错误。为了简化这个过程，PyTorch 提供了一个面向开发者的 API，以确保对 __torch_function__ 覆盖的全面支持。这个 API 是私有的，并且未来可能会在没有警告的情况下进行更改。
首先，要获取所有可覆盖函数的列表，请使用 torch.overrides._get_overridable_functions 。这将返回一个字典，其键是 PyTorch Python API 中的命名空间，其值是该命名空间中可以覆盖的函数列表。例如，让我们打印出可以覆盖的前 5 个函数的名称：
>>> from torch.overrides import get_overridable_functions
>>> func_dict = get_overridable_functions()
>>> nn_funcs = func_dict[torch.nn.functional]
>>> print([f.__name__ for f in nn_funcs[:5])
['adaptive_avg_pool1d', 'adaptive_avg_pool2d', 'adaptive_avg_pool3d',
'adaptive_max_pool1d', 'adaptive_max_pool1d_with_indices']
这个函数列表使得可以遍历所有可重写的函数，然而在实践中，如果不费劲地手动复制每个函数的签名，就无法为这些函数编写测试。为了简化这个过程， torch.overrides._get_testing_overrides 函数返回一个字典，将 PyTorch API 中的可重写函数映射到具有与原始函数相同签名的哑 lambda 函数，这些函数无条件返回 -1。这些函数与 inspect 结合使用时最为有用，可以分析原始 PyTorch 函数的函数签名：
>>> import inspect
>>> from torch.overrides import get_testing_overrides
>>> override_dict = get_testing_overrides()
>>> dummy_add = override_dict[torch.add]
>>> inspect.signature(dummy_add)
<Signature (input, other, out=None)>
最后， torch.overrides.get_ignored_functions 返回一个元组，其中包含 __torch_function__ 明确不能被重写的函数。此列表可以用来确认不在 get_overridable_functions 返回的字典中的函数不能被重写。
扩展 torch 原生 API
虽然 __torch_function__ 允许用户有效地扩展 PyTorch 的纯 Python 组件的行为，但它不允许扩展用 C++ 实现的 PyTorch 部分。为此， Tensor 子类还可以定义 __torch_dispatch__ ，这将能够在 C++ 层面上覆盖行为。
为了有效地使用此功能，了解 PyTorch 原生部分的实现方式很重要。其中最重要的组件就是我们所说的“调度器”（最佳描述可以在这篇博客文章中找到，尽管它有些过时）。正如其名所示，它负责为特定函数调用调用正确的后端函数。例如，在调用 torch.add(a, b) 时，调度器将检查两个参数，确定应该使用哪个“功能”（自动微分、自动类型转换、函数化等）以及哪个“后端”（CPU、CUDA、MPS 等）来处理这个特定调用，并最终调用所有正确的内核。内核经常执行的操作是“重新调度”。例如，当使用自动类型转换在 GPU 上运行您的神经网络时，第一次调用将是自动类型转换内核，它将处理任何潜在的自动类型转换逻辑并重新调度。下一个功能将是自动微分，它将正确创建自动微分图，然后重新调度。最后，我们到达 CUDA 的后端内核，它将启动正确的 CUDA 内核并返回最终结果。 在退出过程中，autograd 会将图附加到输出上，最后 autocast 将有机会在退出时进行任何必要的更新。
分派器的一种配置是所有这些特征和后端键的调用顺序。最新的列表及其顺序可以在 DispatchKey 枚举中的 DispatchKey.h 中找到。为了扩展 torch，本次讨论中重要的排序子集是：
vmap -> Autocast -> Autograd -> ZeroTensor -> Neg/Conj -> Functionalize -> Python -> Backends
对于本次讨论来说，最重要的键是 Python ，因为每个定义了 __torch_dispatch__ 方法的 Tensor 子类都会调用这个功能。用户定义的方法就是从这里被调用，并且行为可以被任意覆盖。从这里再次调用提供的 func 将执行“重新分派”。
该实现的某些重要影响包括：
此代码运行在“所有功能”之下。因此，它仅负责，就像常规后端一样，生成每个张量的输出值（并且可以，也应该忽略所有高级功能，如自动微分、自动类型转换等）。
如果任何高级功能在未重新调度的情况下实现了给定函数，则它将永远不会到达
Python键，因此__torch_dispatch__回调也永远不会被触发。这特别适用于在自动微分级别上不进行重新调度而评估的复合隐式自动微分函数。这是因为复合隐式自动微分函数通过隐式调用其他原生操作来指定其自动微分公式，因此在自动微分级别上，函数被分解为其原生操作，并评估这些操作。在回调到 Python 并包装结果时，使用与常规 PyTorch Python/C++绑定相同的转换。特别是，某些对象无法在 Python 中表示，需要特殊处理（例如，未定义的张量变为 None）。
我们的本机函数以
torch.ops.{namespace}.{func_name}.{overload_name}的形式懒加载为可调用的 Python 对象，以便轻松地从 Python 中与之交互。func对象总是来自此命名空间的一个条目。此命名空间可以直接调用本机操作，绕过常规的 Python API 和绑定代码。
类似地， __torch_function__ 能够介入 torch 的所有 Python API 和 Tensor 方法， __torch_dispatch__ 能够拦截所有对 aten 本机 API 的调用。请注意，在进入分发器之前，Tensor 上的所有方法都转换为函数调用，因此在这里将显示为函数调用： torch.add(a, 2) 和 a + 2 将导致完全相同的 aten 调用。这些函数中的大多数都在 native_functions.yaml 中定义，它指定了这些函数的属性以及它们的后端实现。然后，通过 codegen 自动注册它们的实现和指定功能。一些更特殊的函数或功能也注册在 C++代码库的其他地方或用户定义的 C++扩展中。
也可以使用 torch.library 来添加新的本地函数。这个 Python 特性允许定义和/或添加本地函数的新实现。这可以用来添加缺失的内核、替换现有的内核或定义全新的本地函数。
你可以在子类动物园仓库中找到许多基于 __torch_dispatch__ 的子类的示例。
__torch_dispatch__ 调用约定 ¶
@classmethod
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
pass
当用户使用具有 __torch_dispatch__ 的输入调用运算符时，该调用可能会转发到 __torch_dispatch__ 。在调用 __torch_dispatch__ 之前，args 和 kwargs 会进行归一化，即：
操作符模式中仅包含关键字参数。如果关键字参数等于模式中的默认值，则不会传递。
包含所有其他参数，无论它们是如何传递给操作符的（位置参数或关键字参数）。如果一个参数等于其默认值，并且它是最右边的位置参数或其右侧的所有参数都没有传递，则不会传递。
扩展所有 API 的模式
不幸的是，有一些函数不接受 Tensor 输入。这意味着上述描述的子类方法不能用来覆盖 PyTorch 中所有函数的行为。此外，如果用例需要拦截每个函数调用，将每个 Tensor 改为子类可能会过于侵入。
为了解决这个用例，我们引入了“模式”的概念。这些模式分别用于 __torch_function__ 和 __torch_dispatch__ 覆盖，通过分别子类化 torch.overrides.TorchFunctionMode 和 torch.utils._python_dispatch.TorchDispatchMode 来创建，并用作上下文管理器。
为了简化对它与子类和其他模式交互的描述，每当进入某个模式的上下文管理器时，每个函数的行为都好像在参数列表的开头多了一个 Tensor 参数，该参数为该模式的子类。这意味着特别是所有模式处理程序都将先于任何子类处理程序被调用，并且对应于内部上下文管理器的模式将始终首先运行。
还需要注意的是，在给定的模式处理程序内部，此特定模式被禁用，可以通过执行 with self: 手动重新启用。
下面是一个示例，展示了每种类型的日志模式：
import torch
from torch.overrides import TorchFunctionMode, resolve_name
from torch.utils._python_dispatch import TorchDispatchMode
class FunctionLog(TorchFunctionMode):
def __torch_function__(self, func, types, args, kwargs=None):
print(f"Function Log: {resolve_name(func)}(*{args}, **{kwargs})")
return func(*args, **(kwargs or {}))
class DispatchLog(TorchDispatchMode):
def __torch_dispatch__(self, func, types, args, kwargs=None):
print(f"Dispatch Log: {func}(*{args}, **{kwargs})")
return func(*args, **(kwargs or {}))
def f():
a = torch.rand(10, requires_grad=True)
b = a * 2
b.sum().backward()
print("TorchFunctionMode logging:")
with FunctionLog():
f()
print("TorchDispatchMode logging:")
with DispatchLog():
f()
输出以下内容，并添加额外注释：
TorchFunctionMode logging:
Function Log: torch.rand(*(10,), **{'requires_grad': True})
Function Log: torch.Tensor.mul(*(tensor([0.7164, 0.9897, 0.1745, 0.9336, 0.4287, 0.7989, 0.2169, 0.7474, 0.5624,
0.5970], requires_grad=True), 2), **None)
Function Log: torch.Tensor.sum(*(tensor([1.4328, 1.9794, 0.3490, 1.8671, 0.8573, 1.5977, 0.4338, 1.4948, 1.1249,
1.1939], grad_fn=<MulBackward0>),), **None)
# Note that at the python level, we only see the call to backward but not what happens in the autograd engine.
Function Log: torch.Tensor.backward(*(tensor(12.3307, grad_fn=<SumBackward0>),), **{'gradient': None, 'retain_graph': None, 'create_graph': False, 'inputs': None})
TorchDispatchMode logging:
# Here the requires_grad flag from autograd is removed while default arguments were populated.
Dispatch Log: aten.rand.default(*([10],), **{'device': device(type='cpu'), 'pin_memory': False})
Dispatch Log: aten.mul.Tensor(*(tensor([0.2151, 0.6018, 0.8415, 0.9060, 0.2974, 0.7708, 0.6668, 0.0352, 0.7948,
0.6023], requires_grad=True), 2), **{})
Dispatch Log: aten.sum.default(*(tensor([0.4303, 1.2036, 1.6831, 1.8120, 0.5949, 1.5416, 1.3335, 0.0705, 1.5897,
1.2046], grad_fn=<MulBackward0>),), **{})
# Here we don't see the call to backward itself, but its constituents. Starting here with the factory function that creates the initial gradient.
Dispatch Log: aten.ones_like.default(*(tensor(11.4637, grad_fn=<SumBackward0>),), **{'pin_memory': False, 'memory_format': torch.preserve_format})
# This is the backward of the sum
Dispatch Log: aten.expand.default(*(tensor(1.), [10]), **{})
Dispatch Log: aten.mul.Tensor(*(tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]), 2), **{})
Dispatch Log: aten.detach.default(*(tensor([2., 2., 2., 2., 2., 2., 2., 2., 2., 2.]),), **{})
Dispatch Log: aten.detach.default(*(tensor([2., 2., 2., 2., 2., 2., 2., 2., 2., 2.]),), **{})
