目录
前言
自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考
本次课程我们来学习课程第四章---TensorRT 模型部署优化，一起来学习 sparse tensor core
课程大纲可以看下面的思维导图
0. 简述
本小节目标：理解 NVIDIA 是如何使用 sparse tensor core 来处理带有稀疏性的矩阵乘法
这节课我们来学习剪枝的第三个小部分---sparse tensor core，这个小节我们主要是来理解 NVIDIA 的 sparse tensor core 是如何处理稀疏性的矩阵乘法的，这里做一个科普稍微给大家扩展一下
下面我们开始本次课程的学习🤗
1. 自动驾驶中需要关注的电力消耗
下面是 Jetson AGX Orin 的电力消耗图，它的 Peak performance 和对应的电力消耗是：
- GPU：170 TOPS （Sparse INT8），85 TOPS（DENSE INT8）
- DLA0：52.5 TOPS （Sparse INT8），26.25 TOPS（DENSE INT8）
- DLA1：52.5 TOPS （Sparse INT8），26.25 TOPS（DENSE INT8）
- Power consumption ~60W
Jetson AGX Orin
我们可以看到 Orin 中的 Sparse INT8 的 TOPS 比 Dense INT8 高两倍，那我们自然而然回去想 Sparse 和 Dense 分别代表着什么呢？其实 Sparse 就是 NVIDIA Ampere 架构中所支持的一个东西，下面我们一起来看看
2. Ampere架构中的3rd Generation Tensor core
在 Ampere 架构（A100/Jetson AGX Orin）中的第三代 Tensor core 支持带有 sparsity 的 matrix 计算，更准确来说：
- 支持 fine-grained structured sparsity
- fine-grained 细粒度，主要是对权重的某个元素本身进行分析剪枝，是否归零
- structured 表现在 sparsity 的 pattern 是以 1x4 vector 的大小进行 2:4 的归零（vector-wise pruning），具体如下图所示，每一个 1x4 vector 只保留其中的两个元素即图中绿色部分，另外两个归零即图中白色部分
- 50% 粒度的 sparse pruning，理论上可以实现 2x 的吞吐量的提升
结构化稀疏矩阵
tensor core计算示意图
我们下面来看看 sparse 在 tensor core 中是如何做支持的，假设我们有一个 Dense matrix W 0 W_0 W0，通过 Fine-grained structured pruning 之后变成了一个 structured-sparse matrix W W W，如下图所示：
也就是说一个 dense 的 matrix 会以 2:4 sparsity 的方式进行剪枝（每 4 个连续的权重中最小的两个归零），下图展示了 pruning 的过程：
下面我们以具体的数据来更加形象的说明整个 Fine-grained structured pruning 过程，Original 代表原始的 Dense Matrix，通过剪枝将其中的每四个中的两个进行 pruning 剪枝，得到 2:4 Sparse Matrix：
剪枝完成之后，对于已经 sparse pruning 过的 matrix 可以进一步进行压缩，在 memory 中只保存非零的 weight，至于哪些 weight 是零，哪些 weight 非零，我们可以用一个 2-bits indices 来保存（可以把它理解为一种索引），如下图所示：
这样一来 weight 的大小减半，同时对于 activation values 可以通过 2-bits indices 来决定 activation values 中哪些值是参与计算的，哪些是 skip 掉的（这个过程需要特殊的硬件 unit 来实现），从而实现 2x 的计算吞吐量的提升，整个过程如下图所示：
上图还对比了 Dense Matrix 计算和 Sparse Matrix 计算，可以看到 Sparse Matrix 的计算省去了 Dense Matrix 中很多没必要的计算，它将一个 1x8 与 8x1 的计算通过 indices 变成了一个 1x4 与 4x1 的计算
3. Sparse tensor core做矩阵乘法
我们再来看一个具体的例子，Dence Tensor core（FP16）的计算 A(M,K) * B(K,N) = C(M,N) 的过程可以用下图来表示：
我们知道 A100 中的 tensor core 可以用 1 cycle 完成一个 16x16 * 16x8 = 16 * 8 的矩阵乘法，需要用 2 cycle 完成一个 16x32 * 32x8 = 16 * 8 的矩阵乘法，如下图所示：
上面展示的是 Dense Tensor core（FP16）的计算过程，下面我们来看看 Sparse Tensor core（FP16）的计算 A(M, k) * B(K, N) = C(M, N) 的过程，如下图所示：
如果说 A 中的 matrix 拥有 sparsity，是按照 2:4 的 pattern 进行 pruning，我们可以重构 A，如下图所示，重构后的 A 的 memory 占用空间直接减半达到压缩的效果，我们可以把这里的 A 理解为 conv 或 FC 中的 weight
那么对于 B，可以通过索引对 B 中参与计算的值进行筛选（也可以理解为对 B 的重构），如下图所示，我们可以把这里的 B 理解为 conv 或 FC 中的 activation values。另外这个筛选的过程其实是需要硬件支持的，也就是只能是第三代 tensor core 才能做，其他硬件可能就不行，这也就是为什么 Fine-grained pruning 细粒度剪枝可能会需要特定硬件的支持
这样我们可以得到一个结论：
- 使用 dence tensor core 需要用 2 cycle 完成一个 16x32 * 32x8 = 16 * 8 的矩阵乘法
- 使用 sparse tensor core 只需要 1 cycle 就可以完成一个 16x32 * 32x8 = 16 * 8 的矩阵乘法
然而这里面很容易忽视的一点就是，为了实现 sparse 的计算而添加的额外操作的 overhead
- compress weight 的 overhead
- reconstruct activation values 的 overhead
虽然这些是在硬件上可以完成，从而自动的将 0 参与的计算全部 skip 掉，然而这些多余的操作有时会比较凸显，尤其是当模型并不是很大，参与 sparse 的计算的激活值不是很大时，使用 sparsity 的特性做计算效果不是那么好。目前认为 sparse tensor core 在 NLP 领域的加速可能会比较可观。
上图展示了参与计算的参数量的不同 sparse tensor core 体现出来的加速比也不同，如果参与计算的参数量比较少时，此时 sparse 带来的加速比不是很明显，随着参数量的增加，sparse 带来的加速比也随之增加
总结
本次课程我们主要学习了 NVIDIA 的第三代 tensor core 是如何处理稀疏的矩阵乘法计算的，它主要是通过 Fine-grained pruning 将 Dense Matrix 变成 Sparse Matrix，随后进行 compressed 压缩得到压缩后的 weights 和对应的 indices，通过 indices 可以 skip 掉 weights 和 activation values 之间的一个 0 参与的计算。不过需要注意的是 sparse 的计算会有一些额外的 overhead 导致在参与计算的参数量比较小时加速并不是很可观
OK，以上就是剪枝的第三个小部分 sparse tensor core 的全部内容了，下节我们进入第五章---实战：TensorRT API的基本使用的学习，敬请期待😄