【TVM 教程】在 CPU 上部署 Hugging Face 剪枝模型
内容提要
Apache TVM是一个深度学习编译框架,适用于CPU、GPU和各种机器学习加速芯片。本教程演示了如何使用TVM来加速剪枝后的模型。剪枝是通过将权重值替换为0来减小模型参数大小的技术。修剪结构化稀疏权重的目的是将修剪后的权重聚集在一起,以便跳过整个块。非结构化稀疏权重是根据原始权重值进行修剪的权重。本教程包含结构化和非结构化稀疏。通过使用TVM的块稀疏优化,可以对非结构化的稀疏模型进行加速。
关键要点
-
Apache TVM 是一个深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。
-
本教程演示如何使用剪枝后的模型 PruneBert,并利用 TVM 加速模型稀疏支持。
-
剪枝是一种通过将权重值替换为 0 来减小模型参数大小的技术。
-
修剪结构化稀疏权重的目的是将修剪后的权重聚集在一起,以便跳过整个块。
-
非结构化稀疏权重是根据原始权重值进行修剪的权重,通常在高稀疏下加速效果更明显。
-
本教程包含结构化和非结构化稀疏,PruneBert 模型是非结构化的,具有 95% 的稀疏性。
-
可以用结构化稀疏为未修剪模型生成随机稀疏权重,以比较真实速度与假权重的块稀疏速度。
-
除了 TVM,还需要 scipy 和 TensorFlow(版本在 2.2 以上)。
-
下载和转换 Transformers 模型为 TensorFlow graphdef,并转换为 Relay 计算图。
-
运行导入模型的默认版本,稀疏模型的加速效果依赖于稀疏感知内核的使用。
-
通过应用 Relay pass,将密集层的权重转置并识别可以稀疏的权重矩阵。
-
Block Compressed Row Format (BSR) 用于对张量的非零块进行索引,以提高算法效率。
-
运行 benchmark 测试以评估稀疏模型的加速效果,结果显示可提高约 2.5 倍的速度。