「TVM 教程」在 CPU 上部署 Hugging Face 剪枝模型

作者：HyperAI超神经发布时间：2024-08-05

Hugging Face

更多 TVM 中文文档可访问 →

https://tvm.hyper.ai/docs/how_to/deploy/deploy_models/hugging_face

作者：Josh Fromm

本教程演示如何采用剪枝后的模型（本例中模型是来自 Hugging Face 的 PruneBert），并使用 TVM 来利用模型稀疏支持来加速。

尽管本教程的主要目的是在已经修剪过的模型上实现加速，但评估修剪后模型的速度也十分必要。为此，我们提供了一个函数采用未修剪的模型，并将其权重替换为指定稀疏的随机和修剪权重。确定模型是否值得修剪时，这可能是一个有用的特性。

进入代码前讨论一下稀疏和剪枝，并深入研究两种不同类型的稀疏：结构化和非结构化。

剪枝是一种主要通过将权重值替换为 0 来减小模型参数大小的技术，尽管选择哪些权重设置为 0 的方法众多，但最直接的方法是选择具有最小值的权重。

通常，权重会被修剪为所需的稀疏百分比。例如，一个 95% 的稀疏模型将只有 5% 的权重非零。修剪成非常高的稀疏通常需要微调，或完全重新训练，因为它是有损近似。尽管通过简单的压缩从修剪后的模型中很容易获得参数大小的好处，但利用模型稀疏来产生 runtime 加速更加复杂。

修剪结构化稀疏权重的目的，是把修剪过的权重聚集在一起。换言之，用它们的值和位置进行修剪。将修剪后的权重捆绑在一起的好处是它允许诸如矩阵乘法之类的算法跳过整个块。

事实证明，在当今可用的大多数硬件上，某种程度的块稀疏对于实现显著加速非常重要。这是因为在大多数 CPU 或 GPU 加载内存时，一次跳过读取单个值并不会节省任何工作，而是使用向量化指令之类的东西读入并执行整个块。

非结构化稀疏权重是仅根据原始权重值进行修剪的权重，它们看起来随机分散在整个张量中，而非像块稀疏权重中看到的那样成块。在低稀疏下，非结构化剪枝技术很难加速。然而，在高稀疏下，会出现许多全 0 值的块，这使得加速成为可能。

本教程包含结构化和非结构化稀疏。Hugging Face 的 PruneBert 模型是非结构化的，但 95% 是稀疏的，即使不是最优的，也可以对其应用 TVM 的块稀疏优化。

可以用结构化稀疏为未修剪模型生成随机稀疏权重。将 PruneBert 的真实速度与使用假权重的块稀疏速度比较，可以发现结构化稀疏的优势。

加载所需模块

除了 TVM，还需要 scipy（最新的 transformers）和 TensorFlow（版本在 2.2 以上）。

配置设置

从参数开始，定义要运行的模型类型和稀疏。

下载和转换 Transformers 模型

下面从 transformers 模块获取一个模型，下载并转换为 TensorFlow graphdef，将该 graphdef 转换为可以优化和部署的 Relay 计算图。

转换为 Relay 计算图

目前有很多工具可以获得正确格式的 transformers 模型，从而进行 Relay 转换。下面的函数将导入的计算图保存为 Relay 的 json 格式，这样就不必在每次运行此脚本时从 TensorFlow 重新导入了。

运行密集计算图

运行导入模型的默认版本。注意，即使权重是稀疏的，也不会看到任何加速，因为在这些密集（但大部分为零）张量上使用的是常规密集矩阵乘法，而非稀疏感知内核。

运行稀疏计算图

接下来把计算图转换为稀疏表示，并在需要时生成假的稀疏权重。然后用与 dense 相同的 benchmark 测试脚本来测试速度！对计算图应用一些 Relay pass，从而利用稀疏。

首先用 simple_fc_transpose 将密集层的权重转置到参数中，便于矩阵乘法转换为稀疏版本。接下来应用 bsr_dense.convert 来识别所有可以稀疏的权重矩阵，并自动替换它们。

下面的 bsr_dense.convert 函数通过检查 sparse_threshold 百分比稀疏，识别模型中的哪些权重可以变得稀疏，并将这些权重转换为 Block Compressed Row Format (BSR)。

BSR 本质上是一种对张量的 nonzero chunks 进行索引的表示，使算法可以轻松加载那些 nonzero chunks，并忽略张量的其余部分。一旦稀疏权重采用 BSR 格式，就会应用 relay.transform.DenseToSparse，实际上是用 relay.sparse_dense 函数来替换 relay.dense 操作，从而运行更快。