一、前言
随着深度学习的快速发展，Transformer 算法在各类 AI 任务中都得到了广泛关注。然而，Transformer 的部署面临计算量大、延迟高等挑战，地平线计算芯片能够支撑 Transformer 算法的高效推理，可以成为用户在边缘平台部署 Transformer 算法的优先选择。本文将对 Transformer 性能优化方法做出详细说明。
二、性能优化策略
尽管地平线计算芯片和配套的算法工具链具备高效部署 Transformer 算法的能力，但由于 Transformer 结构的复杂性和特殊性，直接部署公版模型有时会不可避免地会遇到性能瓶颈，导致端侧推理速度较慢。这里介绍几种部署 Transformer 算法时常见的性能问题并提供优化策略。所有算子的耗时分析均基于模型编译时生成的 html 性能预估文件进行。
2.1 优化 Softmax 维度
以 DETR 算法为例，公版算法 17 个 Softmax 的总耗时约为 58ms，优化算法 17 个 Softmax 总耗时约为 42ms，有 16ms 的优化。
这些差异主要体现在 SelfAttention 部分，公版算法是对 1x8x1050x1 的 Shape 做 Softmax，而优化算法是对 8x25x42x1 的 Shape 做 Softmax。这个调整没有改变计算量，但是重排了 Softmax 的计算维度与批次结构，从而提高了并行度与缓存命中率。多个小 Softmax（42 维）比单个大 Softmax（1050 维）更容易并行化，并能让访存更加连续，减少访存延迟。多个 Softmax 同时分配到不同计算单元，使 BPU 计算资源被更高效地利用，从而有效降低计算耗时。
公版算法 Softmax(1x8x1050x1)的静态耗时预估：
地平线优化算法 Softmax(8x25x42x1)的静态耗时预估：
Softmax 算子的量化策略为：将其替换为 6 个等效算子，再逐一对每个算子进行量化。
2.2 优化 BPU 硬件对齐耗时
这里还是以 DETR 算法为例，对于 encoder 中 layers0 的线性层计算，公版算法和地平线优化算法的耗时情况如下：
公版算法
地平线优化算法
可以看到，无论是公版算法还是优化算法，这两层结构的计算量都是相同的，均为 1.1GOPs。但由于输入张量的 Shape 不同，导致耗时差异明显。公版算法的 Shape 是 1x1050x1x2048 和 1x1050x256，受到 BPU 硬件对齐规则限制，第三维从 1 Padding 到 8，大量计算资源消耗在无效数据上，造成了严重的算力浪费。而优化算法的 Shape 是 1x25x42x2048 和 1x25x42x256，Padding 规则仅仅将第二维的 25 变成 26，第三维的 42 变成 48，相比公版 Shape 节约了大量的 BPU 计算资源。仅这一处优化就可降低约 3.6 毫秒耗时，整个模型有多处相似结构，累计的性能提升非常明显。
2.3 优化数据搬运耗时
数据搬运特指 Reshape/Transpose 这类不做数值计算，只做搬运操作的算子。对于 SwinT 算法，公版和优化版本在 Transformer 核心算子（如 Softmax，Layernorm，Matmul）的计算耗时方面表现相当，主要区别在于 Reshape 和 Transpose 的耗时差异，具体如下表所示：
这种优化手段需要用户通过算法设计，尽可能规避模型中出现耗时明显的 Transpose 和 Reshape 算子。但也不是所有的 Transpose 和 Reshape 都有大量耗时，需要结合 html 静态性能报告进一步分析排查。
如果用户选择使用 QAT 链路进行 Transformer 算法的量化部署，算法工具链还提供了更多了性能优化手段。
2.4 使用 QAT 算子避免数据搬运
在 Transformer 结构中，经常需要对 Matmul 的其中一个输入 Transpose 变换维度，这个 Transpose 往往会带来一定的性能开销。用户可以使用 QAT 封装的 Matmul 算子，该算子可以通过入参识别用户希望对哪个输入做维度变换，并在内部做高效处理，从而避免显式的 Transpose 操作。QAT 封装的 Matmul 算子使用示例如下：
Plain
# 原本使用方式
k = k.transpose(-1, -2)
attention = torch.matmul(q, k)
# QAT使用方式
import horizon_plugin_pytorch.nn.quantized as quantized
self.matmul = quantized.FloatFunctional()
attention = self.matmul.matmul(q, k, x_trans=False, y_trans=True)公版的 Layernorm 仅支持对最后若干维度做计算，如果要对中间维度做 Layernorm，需要先 Transpose。而使用 QAT 封装算子可以避免 Transpose，直接对中间维度做 Layernorm。具体使用方式如下：
Plain
# 原本使用方式
# NCHW -> NHWC
x = x.permute(0, 2, 3, 1)
# NHWC 按通道 C 做归一化
self.layernorm = nn.LayerNorm((C))
x = self.layernorm(x)
# NHWC -> NCHW
x = x.permute(0, 3, 1, 2)
# QAT使用方式2.5 使用四维算子代替三维算子
BPU 的硬件特性决定了其对四维数据的支持更加高效，在优化版本的 Transformer 算法中，大量使用了四维算子替换三维算子，以充分发挥 BPU 的计算特性。QAT 也提供了封装好的四维 Transformer 模块供用户使用，如 MultiHeadAttention，PatchMerging 等。
此处展示 QAT 四维算子 HorizonMultiHeadAttention 的部分定义，可在 GPU Docker 中访问/usr/local/lib/python3.10/dist-packageshat/models/base_modules/attention.py 查看完整代码：
Plain
class HorizonMultiheadAttention(nn.Module):
"""modify torch.nn.MultiheadAttention to support quantization.
Args:
embed_dim: Total dimension of the model.
num_heads: Number of parallel attention heads.
Note that ``embed_dim`` will be split across ``num_heads``,
i.e. each head will have dimension ``embed_dim // num_heads``.
dropout: Dropout probability. Default: ``0.0`` (no dropout).
bias: If specified, adds bias to input / output projection layers.DETR 算法的 PatchMerging 结构也可以使用四维算子做替换，示例如下：
Plain
# 原始的使用方式
class PatchMerging(nn.Module):
def _init_(self, dim, norm_layer):
super().__init_()
self.dim = dim
self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
self.norm = norm_layer(4 * dim)
self.cat = quantized.FloatFunctional()
def forward(self,x,H,W):
B, L, C = x.shape2.6 使用 Conv2d 代替 Vector 计算
这里 Vector 并不是指 std::vector 或者一维向量，而是指模型中可以并行处理的一批数据元素，例如一个通道内的像素值。Vector 计算通常指逐元素操作，如均值/求和/取最大值等，与之相对应的是 Tensor 计算。Transformer 算法中比较典型的 Vector 计算有 Elementwise、Reduce 等，模型中如果 Vector 计算占比较多，会影响 BPU 利用率，导致部署性能下降。
在一些情况下，可以使用 Conv2d 等效替换 Vector 操作（如 Mul/Reduce 等）。像是 Conv->ReduceSum->Conv 串接的结构，其中 ReduceSum 就可以替换成 Conv2d：比如 reduce on C 可以构建为 input channel = C，output channel = 1 的 Conv2d；reduce on H/W 可以构建为 kernel_h = H 或 kernel_w = W 的 Conv2d 等。使用 Conv2d 代替 Vector 计算，会提升算法的部署性能，同时不影响计算精度。