PyTorch 2.2 中文官方教程（十三）（2）-阿里云开发者社区

PyTorch 2.2 中文官方教程（十三）（1）https://developer.aliyun.com/article/1482562

获取 PyTorch 运算符的完整列表

PyTorch 提供了一个生成文件build/aten/src/ATen/RegistrationDeclarations.h中的可扩展 C++运算符的完整列表。此文件仅在从源代码构建 PyTorch 后才可用。以下是文件的一部分：

Tensor  abs(const  Tensor  &  self);  // {"schema": "aten::abs(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}
Tensor  &  abs_(Tensor  &  self);  // {"schema": "aten::abs_(Tensor(a!) self) -> Tensor(a!)", "dispatch": "True", "default": "True"}
Tensor  &  abs_out(Tensor  &  out,  const  Tensor  &  self);  // {"schema": "aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "True", "default": "False"}
Tensor  absolute(const  Tensor  &  self);  // {"schema": "aten::absolute(Tensor self) -> Tensor", "dispatch": "False", "default": "False"}
Tensor  &  absolute_(Tensor  &  self);  // {"schema": "aten::absolute_(Tensor(a!) self) -> Tensor(a!)", "dispatch": "False", "default": "False"}
Tensor  &  absolute_out(Tensor  &  out,  const  Tensor  &  self);  // {"schema": "aten::absolute.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "False", "default": "False"}
Tensor  angle(const  Tensor  &  self);  // {"schema": "aten::angle(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}
Tensor  &  angle_out(Tensor  &  out,  const  Tensor  &  self);  // {"schema": "aten::angle.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "True", "default": "False"}
Tensor  sgn(const  Tensor  &  self);  // {"schema": "aten::sgn(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}

与单个运算符相关联的多个字段。让我们以 abs_out 为例进行详细说明：

Tensor & abs_out(Tensor & out, const Tensor & self); 是运算符的 C++ 签名，您的 C++ 内核应该与此签名完全匹配。
aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) 是表示运算符的唯一模式，与 C++ 签名相比，还包含别名和突变注释。这是调度器用来查找运算符的唯一标识符。
dispatch 和 default 是布尔字段，提供了关于原生 PyTorch 内核能够做什么的信息，因此暗示了是否需要后端扩展者实现该内核。更多细节可以在为新后端注册内核中找到。

为新后端注册内核

要将您的内核注册到 PyTorch 调度器中，您可以使用在 C++ 中注册分发运算符中描述的 TORCH_LIBRARY_IMPL API：

TORCH_LIBRARY_IMPL(aten,  PrivateUse1,  m)  {
  m.impl(<schema_my_op1>,  &my_op1);
  m.impl(<schema_my_op2>,  &my_op2);
  m.impl(<schema_my_op2_backward>,  &my_op2_backward);
}

现在让我们深入了解哪些运算符需要来自定制后端的内核以及这些内核的具体内容。

PyTorch 目前有超过 1600 个运算符，而且仍在增长。对于后端扩展来说，跟上这种速度是不现实的。即使对于像 CPU 或 CUDA 这样的原生后端，通常也需要大量工作为每个新运算符编写专用内核。

幸运的是，一些原生 PyTorch 内核是以一种方式编写的，它们分解为几个已知运算符的组合。换句话说，您只需要实现一组已知运算符（下面需要注册的运算符）而不是所有 PyTorch 运算符。

PyTorch 运算符可以分为两类：

需要注册的运算符：这些运算符的 PyTorch 原生实现是特定于后端的，因此需要为定制后端提供内核。否则，在定制后端上调用此类运算符将导致错误。

在 RegistrationDeclarations.h 中，这些运算符在其附带的注释中的元数据中，dispatch 设置为 True 并且 default 设置为 False。

注册是可选的：后端扩展者可以跳过为这些操作注册而不会牺牲任何支持。然而，如果后端扩展者想要覆盖 PyTorch 提供的默认内核，他们仍然可以将他们定制的内核注册到他们的后端，调度器将仅在您的后端中使用它。例如，PyTorch 的 max_pool2d 的当前实现返回 indices 作为前向输出的一部分，这在 torch_xla 中创建了开销，因此 torch_xla 为 max_pool2d 注册了自己的内核。

在 RegistrationDeclarations.h 中，这些运算符在其附带的注释中的元数据中，dispatch 设置为 False 或 default 设置为 True。

新后端的自动求导支持

梯度公式大多是纯数学的，因此对于所有后端都是通用的。PyTorch 经常注册一个用于别名调度键 Autograd 的内核，这意味着它可以被所有后端使用。

对于这些运算符，您不必担心它们的导数公式，您只需在 RegistrationDeclarations.h 中为运算符编写前向定义，PyTorch 将自动为您处理后向。

Tensor  my_op1(const  Tensor&  self,  const  Tensor&  other)  {
  // call your backend-specific APIs to implement my_op so that
  // it matches PyTorch's native behavior
}
TORCH_LIBRARY_IMPL(aten,  PrivateUse1,  m)  {
  m.impl(<schema_my_op1>,  &my_op);
}

在某些情况下，PyTorch 的反向内核实现也是特定于设备的，以便从每个后端中挤出最大性能。对于这些运算符，您将在 RegistrationDeclarations.h 中看到 op_backward 出现为 必需注册。

Tensor  my_op2_backward(const  Tensor&  self,  const  Tensor&  other)  {
  // call your backend-specific APIs to implement my_op2_backward so that
  // it matches PyTorch's native behavior
}
// Note backward kernel is still registered to PrivateUse1 instead of AutogradPrivateUse1.
// PyTorch will wrap your backward kernel with proper autograd setup and then link to it in
// my_op2's AutogradPrivateUse1 kernel.
TORCH_LIBRARY_IMPL(aten,  PrivateUse1,  m)  {
  m.impl(<schema_my_op2>,  &my_op2);
  m.impl(<schema_my_op2_backward>,  &my_op2_backward);
}

在一些罕见情况下，PyTorch 对于某些运算符的梯度公式可能有不适用于所有后端的假设。在这些情况下，后端扩展者可以选择通过将来自 torch::autograd::Function 的内核注册到相应的调度键（例如，如果您的后端使用 PrivateUse1，则为 AutogradPrivateUse1）来覆盖 PyTorch 的 Autograd 层：

class  MyAddFunction  :  public  torch::autograd::Function<MyAddFunction>  {
  public:
  static  Tensor  forward(AutogradContext  *ctx,  torch::Tensor  self,  torch::Tensor  other)  {
  at::AutoNonVariableTypeMode  g;
  return  myadd(self,  other);
  }
  static  tensor_list  backward(AutogradContext  *ctx,  tensor_list  grad_outputs)  {
  auto  grad_output  =  grad_outputs[0];
  return  {grad_output,  grad_output};
  }
};
Tensor  myadd_autograd(const  Tensor&  self,  const  Tensor&  other)  {
  return  MyAddFunction::apply(self,  other)[0];
}
// Register the autograd kernel to AutogradPrivateUse1
TORCH_LIBRARY_IMPL(aten,  AutogradPrivateUse1,  m)  {
  m.impl(<myadd_schema>,  &myadd_autograd);
}
// Register the inference kernel to PrivateUse1
TORCH_LIBRARY_IMPL(aten,  PrivateUse1,  m)  {
  m.impl(<myadd_schema>,  &myadd);
}

通过这种技巧，您可以完全控制后端中my_add运算符的训练和推理行为。这里是pytorch/xla存储库中的一个示例。

构建扩展

通过向 PyTorch 添加 C++扩展来支持外部后端。一旦您准备好内核和注册，您可以通过编写一个使用setuptools编译 C++代码的setup.py脚本来构建 C++扩展。以下是来自pytorch/xla 存储库的一个简化示例：

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
setup(
    name='torch_xla',
    ext_modules=[
        CppExtension(
            '_XLAC',
            torch_xla_sources,
            include_dirs=include_dirs,
            extra_compile_args=extra_compile_args,
            library_dirs=library_dirs,
            extra_link_args=extra_link_args + \
                [make_relative_rpath('torch_xla/lib')],
        ),
    ],
    cmdclass={
        'build_ext': Build,  # Build is a derived class of BuildExtension
    }
    # more configs...
)

有关更多详细信息，请参阅我们的 C++扩展教程。

自定义运算符支持

您的新后端应该与在 Python 中扩展的自定义运算符无缝配合，而无需编写任何新的内核，只要自定义运算符由现有 PyTorch 运算符组成（这些运算符已受到您的后端支持）。

对于在 C++中扩展的自定义运算符，它们通常带有后端特定的 C++内核实现，例如 torchvsion 中的 nms 内核，以及自定义的 Python API，例如 torch.ops.torchvision.nms。为了支持这些运算符，后端扩展者需要为您的后端编写一个 C++内核，并将其正确注册到分发器中的相应命名空间，类似于支持 PyTorch 原生运算符。或者，您还可以在您的扩展中添加一个自定义 API，例如torch_xla.core.functions.nms，以满足这些临时请求。

JIT 支持

正如我们在在 C++中注册分发运算符中提到的，通过 m.impl() API 注册的内核支持以未装箱和装箱方式调用。换句话说，您的定制后端也可以与我们的 JIT 跟踪/脚本前端一起工作，就像树内后端（如 CPU 或 CUDA）一样。您还可以为 JIT 图编写专门的优化传递，但我们不会在这里讨论，因为我们尚未确定 JIT 中的集成点，因此当前后端支持将重点放在急切的前端上。

针对原生 PyTorch 后端进行测试

PyTorch 允许使用其通用设备类型测试框架在多种设备类型上运行测试。您可以在测试如何使用它以及如何添加新设备类型方面找到详细信息。一旦添加，使用通用设备类型测试框架的 PyTorch 测试也将使用您的设备类型运行。查看此 Wiki 页面以了解测试如何实例化的示例。

使用您的设备类型运行 PyTorch 现有的测试套件非常重要，以确保正确性，但并非所有 PyTorch 功能都受到每种设备类型的支持。通用设备类型测试框架允许进行相当大的定制，以便设备类型可以选择运行哪些测试，支持哪些数据类型，甚至在比较张量相等性时使用哪些精度。

使用通用设备类型测试框架并且不随 PyTorch 一起提供的示例设备类型是 XLA。请参阅其对通用设备类型测试框架的扩展，其中包含了测试块列表、数据类型块列表和覆盖测试精度的示例。

通用设备类型测试框架正在积极开发中。要请求功能，请在 PyTorch 的 Github 上提交问题。

向后兼容性

目前，PyTorch 无法保证已注册运算符的向后兼容性。运算符及其模式可能会根据需要进行添加/修改/删除。注册的内核必须与 PyTorch 版本完全相同。如果 PyTorch 为运算符添加更多参数（即使有默认值），您的旧注册将无法工作，直到更新以匹配 PyTorch 的新签名为止。

因此，我们强烈建议独立存储后端扩展器仅与主要的 PyTorch 发布同步，以最大程度地减少开发中的中断。PyTorch 按季度发布。后端扩展器应该加入 #announcement 频道，以获取有关发布的最新更新。

已知问题和其他说明

并非所有测试套件都是设备通用的。可以通过在 PyTorch 代码库中搜索 instantiate_device_type_tests 来找到可扩展的测试类，例如 TestTorchDeviceType, TestViewOps, TestTensorDeviceOps, TestTypePromotion 等。
在 C++ 中没有扩展点用于在自定义后端上序列化 Python 张量对象。目前，您只能通过修改 PyTorch 张量 reduce_ex 方法或在独立存储库中进行 monkey patching 来扩展它。
如果您的后端不允许直接访问内存，则应特别注意支持视图操作，因为它们应该共享存储。对视图张量的更改需要传播到其基张量，反之亦然。
如果您的后端无法与原生 PyTorch 优化器一起使用，则在 C++ 中没有优化器的扩展点，例如需要在向后传递时携带状态以更新像 torch-xla 这样的优化器。目前，这种用例只能通过添加自定义 API 或在独立存储库中进行 monkey patching 来实现。