快捷方式

扩展 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 子类。回想一下,Functions 是 autograd 用来编码操作历史和计算梯度的工具。

本文档的第一部分重点介绍反向模式 AD,因为它是最广泛使用的功能。结尾部分讨论了前向模式 AD 的扩展。

何时使用

通常,如果您想在模型中执行不可微分或依赖于非 PyTorch 库(例如 NumPy)的计算,但仍希望您的操作与其他操作链接并与 autograd 引擎一起工作,请实现自定义函数。

在某些情况下,自定义函数也可以用于提高性能和内存使用率:如果您使用 C++ 扩展 实现了前向和反向传播,您可以将它们包装在 Function 中,以便与 autograd 引擎交互。如果您想减少为反向传播保存的缓冲区数量,可以使用自定义函数将操作组合在一起。

何时不使用

如果您已经可以使用 PyTorch 的内置操作来编写函数,那么它的反向图(很可能)已经能够被 autograd 记录。在这种情况下,您不需要自己实现反向函数。考虑使用普通的 Python 函数。

如果您需要维护状态,即,可训练的参数,您应该(也)使用自定义模块。有关扩展 torch.nn 的更多信息,请参阅下面的部分。

如果您想在反向传播期间更改梯度或执行副作用,请考虑注册 张量模块 钩子。

如何使用

采取以下步骤: 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: 您有责任正确使用 ctx 中的函数,以确保新的 Function 与 autograd 引擎正常工作。

  • save_for_backward() 必须用于保存任何要在反向传播中使用的张量。非张量应直接存储在 ctx 上。如果为反向传播保存了既不是输入也不是输出的张量,您的 Function 可能不支持双重反向传播(请参阅步骤 3)。

  • mark_dirty() 必须用于标记任何被前向函数就地修改的输入。

  • mark_non_differentiable() 必须用于告知引擎输出是否不可微分。默认情况下,所有可微分类型的输出张量都将被设置为需要梯度。不可微分类型(即,整数类型)的张量永远不会被标记为需要梯度。

  • set_materialize_grads() 可用于告知 autograd 引擎优化梯度计算,在输出不依赖于输入的情况下,通过不具体化传递给反向函数的梯度张量。也就是说,如果设置为 False,Python 中的 None 对象或 C++ 中的“未定义张量”(x.defined() 为 False 的张量 x)将不会在调用反向传播之前转换为填充零的张量,因此您的代码需要将此类对象视为填充零的张量。此设置的默认值为 True。

步骤 3: 如果您的 Function 不支持双重反向传播,您应该通过使用 once_differentiable() 修饰器来显式声明这一点。使用此修饰器,尝试通过您的函数执行双重反向传播将产生错误。有关双重反向传播的更多信息,请参阅我们的双重反向传播教程。

步骤 4: 建议您使用 torch.autograd.gradcheck() 检查您的反向函数是否通过使用您的反向函数计算雅可比矩阵,并将该值与使用有限差分数值计算的雅可比矩阵逐元素比较,从而正确计算了前向传播的梯度。

示例

下面您可以找到 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)

在这里,我们给出了一个由非张量参数参数化的函数的附加示例

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() 中计算的任何“中间”张量被保存,则它们必须作为输出返回,或者组合 forwardsetup_context() (参见 组合或单独的 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 有两种主要方法。要么

  • 定义一个 forward(),它将前向计算逻辑与 setup_context() 结合起来

  • (从 PyTorch 2.0 开始)定义一个单独的 forward()setup_context()

我们推荐第二种选择(单独的 forward()setup_context()),因为这更接近 PyTorch 原生操作的实现方式,并且可以与 torch.func 变换组合。但是,我们计划在未来支持这两种方法;将 forward()setup_context() 结合:导致更大的灵活性,因为您能够保存中间值而无需将它们作为输出返回。

请参阅上一节,了解如何定义具有单独 forward()setup_context()Function

这是一个如何定义具有组合 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() 函数必须与 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 命名空间中接受 Tensor 操作数的函数(如 torch.add()),该怎么办?

如果您的自定义 Python 类型定义了一个名为 __torch_function__ 的方法,则当您的自定义类的实例传递给 torch 命名空间中的函数时,PyTorch 将调用您的 __torch_function__ 实现。这使得可以为 torch 命名空间中的任何函数定义自定义实现,您的 __torch_function__ 实现可以调用这些函数,从而允许您的用户将您的自定义类型与他们已经为 Tensor 编写的现有 PyTorch 工作流程一起使用。这适用于与 Tensor 无关的“鸭子类型”以及 Tensor 的用户定义子类。

使用类似 Tensor 的类型扩展 torch

注意

此功能受 NumPy __array_function__ 协议的启发。有关更多详细信息,请参阅 NumPy 文档NEP-0018

为了使之具体化,让我们从一个简单的示例开始,该示例说明了 API 分发机制。我们将创建一个自定义类型,该类型表示 2D 标量张量,由阶数 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

ScalarTensor 添加 __torch_function__ 实现使得上述操作可以成功。让我们重新实现我们的代码,这次添加 __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__ 的类张量类型的列表;args,传递给函数的参数元组;以及 kwargs,传递给函数的关键字参数字典。它使用名为 HANDLED_FUNCTIONS 的全局分发表来存储自定义实现。此字典的键是 torch 命名空间中的函数,值是 ScalarTensor 的实现。

注意

使用全局分发表不是 __torch_function__ API 的强制部分,它只是用于构建覆盖实现的有用设计模式。

此类的定义还不足以使 torch.mean 在我们向其传递 ScalarTensor 时执行正确的操作 – 我们还需要为 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.meanScalarTensor 一起使用

>>> 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 实现不接受 alphaout 作为关键字参数,就像 torch.add() 一样

>>> 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 函数。例如,如果我们更改 ScalarTensor__torch_function__ 实现,如下所示

@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() 将正常工作,尽管即使两个操作数都是 ScalarTensor 实例,返回类型也始终是 Tensor 而不是 ScalarTensor

>>> s = ScalarTensor(2, 2)
>>> torch.mul(s, s)
tensor([[4., 0.],
        [0., 4.]])

另请参阅下面的 MetadataTensor 示例,了解此模式的另一种变体,但始终返回 MetadataTensor 以通过 torch API 中的操作传播元数据。

__torch_function__ 协议旨在完全覆盖 API,部分覆盖可能会导致不良结果,特别是某些函数引发 TypeError。对于子类来说尤其如此,即使它们返回完全相同的结果,也必须覆盖 torch.addtorch.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,它将元数据字典附加到 Tensor,该字典通过 torch 操作传播。由于这是对整个 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 中的命名空间,其值是该命名空间中可以覆盖的函数列表。例如,让我们打印 torch.nn.functional 中前 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 内核并返回最终结果。在退出时,自动梯度会将图附加到输出,最后,自动混合精度将有机会在退出时执行任何需要的更新。

调度器的一种配置是调用所有这些功能和后端键的顺序。最新的列表及其顺序可以在 DispatchKey.h 内的 DispatchKey 枚举中找到。为了扩展 torch 的目的,此讨论的排序的重要子集是

vmap -> Autocast -> Autograd -> ZeroTensor -> Neg/Conj -> Functionalize -> Python -> Backends

对于本次讨论,最重要的键是 Python,因为每个定义了 __torch_dispatch__ 方法的 Tensor 子类都将调用此功能。正是从那里调用了用户定义的方法,并且可以在那里任意覆盖行为。从那里,再次调用提供的 func 将执行“重新分发”。

此实现的一些重要含义是

  • 此代码在“所有功能之下”运行。因此,它仅负责(像常规后端一样)生成每个 Tensor 的输出值(并且可以并且应该忽略所有高级功能,如自动梯度、自动混合精度等)。

  • 如果任何高级功能在不重新分发的情况下实现给定的函数,则它永远不会到达 Python 键,因此永远不会触发 __torch_dispatch__ 回调。对于 CompositeImplicitAutograd 函数来说尤其如此,这些函数在 Autograd 级别进行评估,而无需重新分发。这是因为 CompositeImplicitAutograd 函数通过隐式调用其他原生操作来指定其自动梯度公式,因此在 Autograd 级别,该函数被分解为原生操作,并评估这些原生操作。

  • 当回调到 Python 并且包装结果时,使用的转换与常规 PyTorch Python/C++ 绑定相同。特别是,某些对象无法在 Python 中表示,需要特殊处理(例如,未定义的 Tensor 变为 None)。

  • 我们的原生函数被延迟填充为 torch.ops.{namespace}.{func_name}.{overload_name} 作为可调用的 Python 对象,以便能够轻松地从 Python 与它们交互。提供给 __torch_dispatch__func 对象始终是此命名空间中的条目。此命名空间可用于直接调用原生操作并绕过常用的 Python API 和绑定代码。

__torch_function__ 能够介入所有 torch 的 Python API 和 Tensor 方法的方式类似,__torch_dispatch__ 能够拦截对 aten 原生 API 的所有调用。请注意,Tensor 上的所有方法在进入调度器之前都会转换为函数调用,因此将在此处显示为函数调用:torch.add(a, 2)a + 2 将导致完全相同的 aten 调用。这些函数中的大多数在 native_functions.yaml 中定义,该文件指定了这些函数的属性以及它们的后端实现。它们的实现以及指定的功能然后通过代码生成自动注册。一些更奇特的功能或特性也在 C++ 代码库或其他用户定义的 C++ 扩展中的其他位置注册。

也可以使用 torch.library 添加 新的 原生函数。此 Python 功能允许定义和/或向原生函数添加新的实现。这可用于添加缺少的内核、替换现有内核或定义全新的原生函数。

您可以在 子类动物园 仓库中找到许多基于 __torch_dispatch__ 的子类的示例。

__torch_dispatch__ 调用约定

@classmethod
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
    pass

当用户使用具有 __torch_dispatch__ 的输入调用运算符时,该调用可能会转发到 __torch_dispatch__。args 和 kwargs 在调用 __torch_dispatch__ 之前会被规范化,即

  • kwargs 由运算符模式中的仅关键字参数组成。如果 kwarg 等于其默认值(在模式中),则不会传递它。

  • args 由所有其他参数组成,无论它们是如何传递给运算符的(位置参数与关键字参数)。如果 arg 等于其默认值,并且它是最右边的位置参数,或者其右边的所有参数都未传递,则不会传递它。

使用模式扩展所有 torch API

不幸的是,有些函数不接受 Tensor 输入。这意味着上面描述的子类方法不能用于覆盖 PyTorch 的所有函数的行为。此外,如果用例需要拦截每个函数调用,则将每个 Tensor 更改为子类可能会过于侵入性。

为了解决此用例,我们引入了“模式”的概念。这些模式存在于 __torch_function____torch_dispatch__ 覆盖中,分别通过子类化 torch.overrides.TorchFunctionModetorch.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.]),), **{})

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

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

查看资源