PyTorch’s Tracing Based Selective Build

由 Dhruv Matani，Suraj Subramanian 撰写

引言

TL;DR：在移动设备、SBC（单板计算机）和 IOT 设备上运行 PyTorch 可能具有挑战性。当编译时，PyTorch 库很大，包括可能对设备上使用不必要的一些依赖项。

要在设备上运行特定的模型集，我们实际上只需要 PyTorch 库中的一个小子集的功能。我们发现，使用通过选择性构建生成的 PyTorch 运行时，可以实现对二进制大小的 90%以上的减少（在 Linux 上的 x86-64 构建的 CPU 和 QuantizedCPU 后端）。在这篇博客中，我们分享了使用选择性构建生成特定模型的最小运行时的经验，并展示了如何做到这一点。

这对应用开发者来说为什么很重要？

使用由选择性构建生成的 PyTorch 运行时可以减少 AI 应用程序的大小 30+ MB - 对于一个典型的移动应用程序来说这是一个显著的减少！使移动应用程序更轻量级有许多好处 - 它们可以在更广泛的设备上运行，消耗更少的蜂窝数据，并且可以在用户的设备上更快地下载和更新。

开发者体验是什么样的？

此方法可以无缝地与任何现有的 PyTorch Mobile 部署工作流程协同工作。您需要做的所有事情就是用为您的应用程序中要使用的特定模型定制的运行时替换通用的 PyTorch 运行时库。此过程的一般步骤是：

以仪器模式构建 PyTorch 运行时（这被称为 PyTorch 的仪器构建）。这将记录使用的算子、内核和功能。
使用提供的 model_tracer 二进制文件运行此仪表化构建，这将生成一个存储模型所用所有特征的单一 YAML 文件。这些特征将保留在最小运行时中。
使用此 YAML 文件作为输入构建 PyTorch。这是一种选择性构建技术，它可以大大减少最终 PyTorch 二进制文件的大小。
使用此选择性构建的 PyTorch 库来减小您的移动应用程序的大小！

以特殊的“仪表化”模式（通过传递 TRACING_BASED=1 构建选项）构建 PyTorch 运行时，将生成 PyTorch 的仪表化构建运行时以及 model_tracer 二进制文件。使用此构建运行模型可以追踪模型使用的 PyTorch 部分。

图 1：PyTorch 的仪器构建

# Clone the PyTorch repo
git clone https://github.com/pytorch/pytorch.git
cd pytorch

# Build the model_tracer
USE_NUMPY=0 USE_DISTRIBUTED=0 USE_CUDA=0 TRACING_BASED=1 \
  python setup.py develop

现在这个仪器构建被用来运行具有代表性输入的模型推理。模型追踪器二进制文件观察了推理运行期间激活的仪器构建部分，并将其输出到 YAML 文件中。

图 2：在仪器构建上运行模型（们）生成的 YAML 文件

# Generate YAML file
./build/bin/model_tracer \
  --model_input_path /tmp/path_to_model.ptl \
  --build_yaml_path /tmp/selected_ops.yaml

现在我们再次构建 PyTorch 运行时，但这次使用追踪器生成的 YAML 文件。现在的运行时只包含需要此模型的部分。下面图表中称之为“选择性构建的 PyTorch 运行时”。

# Clean out cached configuration
make clean

# Build PyTorch using Selected Operators (from the YAML file)
# using the host toolchain, and use this generated library
BUILD_PYTORCH_MOBILE_WITH_HOST_TOOLCHAIN=1 \
USE_LIGHTWEIGHT_DISPATCH=0 \
BUILD_LITE_INTERPRETER=1 \
SELECTED_OP_LIST=/tmp/selected_ops.yaml \
TRACING_BASED=1 \
  ./scripts/build_mobile.sh

图 3：在选择性构建的 PyTorch 运行时上选择性地构建 PyTorch 和模型执行

给我看代码！

我们整理了一个笔记本，用简单的 PyTorch 模型演示上述过程在代码中的样子。

对于在 Android/iOS 上部署此教程的更实际教程，本教程应该很有帮助。

技术常见问题解答

为什么 PyTorch 的选型构建需要追踪？

在 PyTorch 中，CPU 内核可以通过 PyTorch 调度器调用其他算子。仅仅包括模型直接调用的根算子集合是不够的，因为可能还有许多其他算子被间接调用。在具有代表性的输入上运行模型并观察实际调用的算子列表（即“追踪”）是确定 PyTorch 哪些部分被使用的最准确的方法。

此外，诸如内核应处理哪些数据类型等因素也是依赖于模型提供的实际输入的运行时特性。因此，追踪机制非常适合此目的。

使用基于跟踪的选建功能可以选择哪些功能（包含或排除）？

在基于跟踪的选建过程中，可以选定的 PyTorch 运行时功能如下：

PyTorch 的 ATen 运算符的 CPU/量化 CPU 内核：如果一个 PyTorch 运算符不需要针对选建运行时构建的模型，则该 CPU 内核的注册在运行时中被省略。这通过 Torchgen 代码生成器进行控制。
主要运算符：这由名为 TORCH_SELECTIVE_SCHEMA 的宏控制（通过模板化选建），根据生成的头文件中的信息选择或排除主要运算符。
处理 CPU 内核中特定数据类型的代码：这是通过在由宏 AT_PRIVATE_CHECK_SELECTIVE_BUILD 生成的 switch case 语句中的特定 case 语句中生成异常抛出来实现的。
注册扩展 PyTorch 的自定义 C++类：这由宏 TORCH_SELECTIVE_CLASS 控制，当注册自定义 C++类时可以使用。应与宏 TORCH_SELECTIVE_CLASS 一起使用 torch::selective_class_<>辅助器。

构建过程中使用的 YAML 文件的结构是什么？

追踪后生成的 YAML 文件如下所示。它编码了上述“可选择的”构建功能的所有元素。

include_all_non_op_selectives: false
build_features: []
operators:
    aten::add.Tensor:
        is_used_for_training: false
        is_root_operator: true
        include_all_overloads: false
    aten::len.t:
        is_used_for_training: false
        is_root_operator: true
        include_all_overloads: false
kernel_metadata:
    _local_scalar_dense_cpu:
    - Float
    add_stub:
    - Float
    copy_:
    - Bool
    - Byte
    mul_cpu:
    - Float
custom_classes: []

代码究竟是如何从生成的二进制文件中消除的？

根据具体场景，有 2 种主要技术被用来提示编译器和链接器关于未使用和不可达的代码。然后编译器或链接器会将其作为不可达代码进行清理。

[1] 链接器移除未引用的函数

当一个函数没有被任何可见函数间接引用，并且存在于正在链接的编译对象文件中时，如果提供了正确的构建标志，链接器将会移除它。这种技术通过选择性的构建系统在 2 个场景中被利用。

分发器中的内核注册

如果操作员的内核不需要，则不会与分发器注册。未注册的内核意味着该函数不可访问，它将被链接器删除。

模板化选择性构建

这里的基本思想是使用类模板特化来选择一个类，该类要么捕获对函数的引用，要么不捕获（取决于是否使用），然后链接器可以来清理未引用的函数。

例如，在下面的代码中，没有对函数“ fn2 ”的引用，因此链接器会将其清理掉，因为它在任何地方都没有被引用。

#include <vector>
#include <cstdio>

template <typename T, bool>
struct FunctionSelector {
    T fn_;
    FunctionSelector(T fn): fn_(fn) {}
    T get() { return this->fn_; }
};

// The "false" specialization of this class does NOT retain the argument passed
// to the class constructor, which means that the function pointer passed in
// is considered to be unreferenced in the program (unless it is referenced
// elsewhere).
template <typename T>
struct FunctionSelector<T, false> {
    FunctionSelector(T) {}
};

template <typename T>
FunctionSelector<T, true> make_function_selector_true(T fn) {
    return FunctionSelector<T, true>(fn);
}

template <typename T>
FunctionSelector<T, false> make_function_selector_false(T fn) {
    return FunctionSelector<T, false>(fn);
}

typedef void(*fn_ptr_type)();

std::vector<fn_ptr_type> fns;

template <typename T>
void add_fn(FunctionSelector<T, true> fs) {
    fns.push_back(fs.get());
}

template <typename T>
void add_fn(FunctionSelector<T, false>) {
    // Do nothing.
}

// fn1 will be kept by the linker since it is added to the vector "fns" at
// runtime.
void fn1() {
    printf("fn1\n");
}

// fn2 will be removed by the linker since it isn't referenced at all.
void fn2() {
    printf("fn2\n");
}

int main() {
    add_fn(make_function_selector_true(fn1));
    add_fn(make_function_selector_false(fn2));
}

[2] 编译器消除的无效代码

C++ 编译器可以通过静态分析代码的控制流来检测无效（不可达）的代码。例如，如果存在一个在无条件异常抛出之后的代码路径，那么它之后的所有代码都将被标记为无效代码，并且不会被编译器转换为目标代码。通常，编译器需要使用 -fdce 标志来消除无效代码。

在下面的示例中，您可以看到左侧的 C++ 代码（在红色框中）在右侧没有对应的生成目标代码。

图 4：C++编译器的死代码消除

该属性被用于 PyTorch 内核实现体中，这些实现体包含大量重复代码以处理 Tensor 的多个数据类型。数据类型是 Tensor 存储元素的基本数据类型。这可以是 float、double、int64、bool、int8 等之一。

几乎所有的 PyTorch CPU 内核都使用 AT_DISPATCH_ALL_TYPES*形式的宏，该宏用于为内核需要处理的每个数据类型替换一些代码。例如：

AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(
    kBool, kHalf, kBFloat16, dtype, "copy_kernel", [&] {
  cpu_kernel_vec(
      iter,
      [=](scalar_t a) -> scalar_t { return a; },
      [=](Vectorized<scalar_t> a) -> Vectorized<scalar_t> { return a; });
});

宏 AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3 内部有一个类似于图 4 中代码的 switch-case 语句。跟踪过程记录了触发内核标签“ copy_kernel ”的数据类型，构建过程处理这些标签，并在每个处理不需要为此内核标签的数据类型的 case 语句中插入 throw 语句。

这是如何在 PyTorch 的基于跟踪的选型构建中实现 dtype 选择性的。

结论

基于跟踪的选型构建是一种实用且可扩展的方法，只选择应用程序中使用的部分以保留静态分析无法检测到的代码。这种代码在本质上通常非常依赖于数据/输入。

本文详细介绍了基于跟踪的选型构建在底层的工作原理及其实现的相关技术细节。这些技术也可以应用于其他可以从减小二进制大小中受益的应用和场景。

PyTorch 基于跟踪的选建

引言