PyTorch 2.2 中文官方教程（十三）（3）

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

通过 PrivateUse1 促进新后端集成

原文：pytorch.org/tutorials/advanced/privateuseone.html

译者：飞龙

协议：CC BY-NC-SA 4.0

在本教程中，我们将逐步介绍通过PrivateUse1将存放在pytorch/pytorch存储库之外的新后端集成的一些必要步骤。请注意，本教程假定您已经对 PyTorch 有基本的了解，您是 PyTorch 的高级用户。

注意

本教程仅涉及与 PrivateUse1 机制相关的部分，以促进新设备的集成，其他部分将不予涵盖。同时，并非所有本教程涉及的模块都是必需的，您可以根据实际需求选择对您有帮助的模块。

PrivateUse1 是什么？

在 Pytorch 2.0 之前，PyTorch 为原型外后端扩展提供了三个保留的调度键（及其对应的 Autograd 键），这三个调度键如下：

• PrivateUse1/AutogradPrivateUse1
  
• PrivateUse2/AutogradPrivateUse2
  
• PrivateUse3/AutogradPrivateUse3
  

原型验证通过后，可以申请新后端的私钥，如 CUDA、XLA、MPS 等。

然而，随着 PyTorch 的快速发展，越来越多的硬件制造商尝试将他们的后端集成到 PyTorch 中，这可能会引发以下问题：

• 每个新后端集成都涉及大量文件修改
  
• 目前对调度键数量（DispatchKeySet 64 位限制）有硬性限制
  

注意

通过 PrivateUse1 Key 将新后端集成到 PyTorch 中也存在问题，因为不可能同时集成多个后端。幸运的是，这些原型外后端很少同时使用。

鉴于上述原因，社区开始建议通过PrivateUse1将新后端集成到 PyTorch 中。

然而，之前的PrivateUse1机制并不能完全与新后端集成，因为在某些模块中缺乏相关支持，如 Storage、AMP、Distributed 等。

随着 Pytorch 2.1.0 的到来，针对PrivateUse1的一系列优化和增强已经针对新后端集成进行了，现在可以快速高效地支持新设备的集成。

如何通过 PrivateUse1 集成新后端

在本节中，我们将讨论通过PrivateUse1将新后端集成到 Pytorch 中的细节，主要包括以下部分：

1. 为新后端注册内核。
   
2. 为新后端注册生成器。
   
3. 为新后端注册设备保护。
   
4. 为新后端元数据注册序列化和反序列化函数。
   
5. 其他模块。
   

为新后端注册内核

新后端可能具有一些高性能的运算符实现，可以通过TORCH_LIBRARY_IMPL API 在 Registering a Dispatched Operator in C++中描述的方式注册到调度程序。这涉及几种情况：

1. 为新后端支持的所有前向运算符注册到调度程序，并同时注册回退，以便当新后端不支持某些运算符时，这些运算符可以回退到 CPU 执行，以确保功能的可用性。

at::Tensor  wrapper_Custom_Tensor_add(const  at::Tensor  &  self,  const  at::Tensor  &  other,  const  at::Scalar  &  alpha)  {
  // Implementation of add kernel in new backend
  ...
}
TORCH_LIBRARY_IMPL(aten,  PrivateUse1,  m)  {
  ...
  m.impl("add.Tensor",  TORCH_FN(wrapper_Custom_Tensor_add));
  ...
}
void  custom_cpu_fallback(const  c10::OperatorHandle&  op,  torch::jit::Stack*  stack)  {
  // Add some hints about new devices that do not support and need to fall back to cpu
  at::native::cpu_fallback(op,  stack);
}
TORCH_LIBRARY_IMPL(_,  PrivateUse1,  m)  {
  m.fallback(torch::CppFunction::makeFromBoxedFunction<&custom_cpu_fallback>());
} 

1. 如果新后端需要覆盖PyTorch Autograd layer，则通过AutogradPrivateUse1将torch::autograd::Function的内核注册到调度程序，调度程序和自动求导系统将自动调用这些运算符的前向和后向实现。

class  CumtomSeluFunction  :  public  torch::autograd::Function<CumtomSeluFunction>  {
  // Implementation of selu kernel in new backend
}
at::Tensor  wrapper_AutogradCumstom__selu(const  at::Tensor  &  self)  {
  return  CumtomSeluFunction::apply(self);
}
TORCH_LIBRARY_IMPL(aten,  AutogradPrivateUse1,  m)  {
  ...
  m.impl("selu",  TORCH_FN(wrapper_AutogradCustom__selu));
  ...
} 

1. 通过AutocastPrivateUse1将想要支持自动混合精度（AMP）和回退机制的内核注册到调度程序，当需要时，自动转换系统将自动调用这些内核。

TORCH_LIBRARY_IMPL(aten,  AutocastPrivateUse1,  m)  {
  ...
  KERNEL_PRIVATEUSEONE(<operator>,  <policy>)
  ...
}
TORCH_LIBRARY_IMPL(_,  AutocastPrivateUse1,  m)  {
  m.fallback(torch::CppFunction::makeFallthrough());
} 

需要补充的是，如果要在新后端支持 AMP，需要通过torch._register_device_module("backend_name", BackendModule)注册一个新的BackendModule，并且BackendModule需要具有以下 API：

• get_amp_supported_dtype() -> List[torch.dtype]
  在 AMP 中获取新后端支持的dtype，可能支持一个以上的dtype。
  
• is_autocast_enabled() -> bool
  检查新后端是否启用 AMP。
  
• get_autocast_dtype() -> torch.dtype
  在 AMP 中获取新后端支持的dtype，该dtype由set_autocast_dtype或默认dtype设置，而默认dtype为torch.float16。
  
• set_autocast_enabled(bool) -> None
  在新后端上启用或禁用 AMP。
  
• set_autocast_dtype(dtype) -> None
  在 AMP 中设置新后端支持的dtype，并且dtype包含在从get_amp_supported_dtype获取的dtypes中。
  

为新后端注册生成器

需要支持与新设备对应的生成器。目前，PrivateUse1可以动态注册自定义生成器，主要分为以下几个步骤。

1. 继承GeneratorImpl类以实现与新后端对应的生成器类，并实现各种通用方法。
   
2. 定义一个带有单个参数device index的新后端builder。
   
3. 调用REGISTER_GENERATOR_PRIVATEUSE1宏完成动态注册。
   

struct  CustomGeneratorImpl  :  public  c10::GeneratorImpl  {
  // Implementation of generator in new backend
}
at::Generator  make_custom_generator(c10::DeviceIndex  device_index)  {
  return  at::make_generator<CustomGeneratorImpl>(device_index);
}
REGISTER_GENERATOR_PRIVATEUSE1(make_cumstom_generator) 

为新后端注册设备保护

PyTorch 通过DeviceGuard提供了与设备、流和事件切换相关的功能。这个功能也适用于PrivateUse1关键。

1. 继承DeviceGuardImplInterface类以实现与新后端对应的各种通用方法。
   
2. 调用C10_REGISTER_GUARD_IMPL宏完成动态注册。
   

struct  CustomGuardImpl  final  :  public  c10::impl::DeviceGuardImplInterface  {
  // Implementation of guard in new backend
}
C10_REGISTER_GUARD_IMPL(PrivateUse1,  CustomGuardImpl); 

为新后端元数据注册序列化和反序列化函数

PyTorch 目前能够动态注册序列化/反序列化函数，以支持在TensorImpl.ExtraMeta类中命名为backend_meta_的新后端附加元数据的序列化和反序列化。您可以参考以下步骤：

1. 继承BackendMeta类以实现与新后端对应的CustomBackendMetadata，并且新后端的各个字段可以在类中自定义。
   
2. 实现新后端的序列化和反序列化函数，函数签名为void(const at::Tensor&, std::unordered_map&)。
   
3. 调用TensorBackendMetaRegistry宏完成动态注册。
   

struct  CustomBackendMetadata  :  public  c10::BackendMeta  {
  // Implementation of backend metadata in new backend
}
void  for_serialization(const  at::Tensor&  t,  std::unordered_map<std::string,  bool>&  m)  {
  // Implementation of serialization
}
void  for_deserialization(const  at::Tensor&  t,  std::unordered_map<std::string,  bool>&  m)  {
  // Implementation of deserialization
}
TensorBackendMetaRegistry(c10::DeviceType::PrivateUse1,  &for_serialization,  &for_deserialization); 

其他模块

除了上述部分外，还有一些其他模块可以通过PrivateUse1进行扩展，例如分布式集体通信、基准计时器等，这些将在未来添加。关于PrivateUse1集成的一个示例是Ascend NPU。

如何通过 Privateuse1 改进用户体验

通过PrivateUse1集成新设备的主要目标是满足基本的功能要求，接下来要做的是改进可用性，主要涉及以下几个方面。

1. 向 PyTorch 注册新的后端模块。
   
2. 生成与新后端相关的方法和属性。
   
3. 生成与新后端相关的方法和属性。
   

向 PyTorch 注册新的后端模块

PyTorch 中的一些与 CUDA 相关的接口可以通过以下形式调用：torch.cuda.xxx。因此，为了符合用户习惯，通过PrivateUse1机制实现的新后端也应该提供类似的接口。

例如，使用Ascend NPU：

torch._register_device_module('npu', torch_npu.npu) 

完成上述操作后，用户可以通过torch.npu.xxx调用Ascend NPU的一些独有 API。

将 PrivateUse1 重命名为新后端的自定义名称

PrivateUse1 键是集成到 PyTorch 中的新后端的内部机制。对于用户来说，与 PrivateUse1 相比，与新后端密切相关的自定义名称应该更加友好。

以 Ascend NPU 为例，第一种用法将更加用户友好。

torch.rand((2,2),device='npu:0')
torch.rand((2,2),device='privateuse1:0') 

现在，PyTorch 为自命名的 PrivateUse1 后端提供了一个新的 C++/Python API，非常简单易用。

torch.rename_privateuse1_backend("npu") 

c10::register_privateuse1_backend("npu") 

生成与新后端相关的方法和属性

将 PrivateUse1 重命名为自定义名称后，在新后端的 Tensor, nn, Storage 模块中自动生成与新后端名称相关的属性和方法。

这里以 Ascend NPU 为例：

torch.rename_privateuse1_backend("npu")
unsupported_dtype = [torch.quint8]
torch.utils.generate_methods_for_privateuse1_backend(for_tensor=True, for_module=True, for_storage=True, unsupported_dtype=unsupported_dtype) 

然后，您可以使用以下方法和属性：

torch.Tensor.npu()
torch.Tensor.is_npu
torch.Storage.npu()
torch.Storage.is_npu
...

未来工作

PrivateUse1 机制的改进仍在进行中，因此新模块的 PrivateUse1 集成方法将逐步添加。以下是我们正在积极开展的几个项目：

• 添加 分布式集体通信 的集成方法。
  
• 添加 基准计时器 的集成方法。
  

结论

本教程指导您通过 PrivateUse1 将新后端集成到 PyTorch 中的过程，包括但不限于运算符注册、生成器注册、设备保护注册等。同时，介绍了一些方法来改善用户体验。

模型优化

使用 TensorBoard 的 PyTorch 分析器

原文：pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html

译者：飞龙

协议：CC BY-NC-SA 4.0

注意

点击这里下载完整示例代码

本教程演示了如何使用 TensorBoard 插件与 PyTorch 分析器来检测模型的性能瓶颈。

介绍

PyTorch 1.8 包括一个更新的分析器 API，能够记录 CPU 端操作以及 GPU 端的 CUDA 内核启动。分析器可以在 TensorBoard 插件中可视化这些信息，并提供性能瓶颈的分析。

在本教程中，我们将使用一个简单的 Resnet 模型来演示如何使用 TensorBoard 插件来分析模型性能。

设置

要安装torch和torchvision，请使用以下命令：

pip install torch torchvision 

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