SwinTransformer模型优化

2023/08/04 18:07:03

1.SwinTransformer概述#

自从Transformer在NLP任务上取得突破性的进展之后，业内一直尝试着把Transformer用于CV领域。之前的若干尝试都是将Transformer用在了图像分类领域，但这些方法都面临两个非常严峻的挑战，一是多尺度问题，二是计算复杂度的问题。

基于这两个挑战，swint的作者提出了一种层级式提取的Transformer，并通过移动窗口的方式来学习特征。在窗口内计算自注意力可以带来更高的效率；同时通过移动的操作，让相邻的窗口之间有了交互，变相达到了一种全局建模的能力，进而解决了上面两个问题。

Swin Transformer将transformer结构与cnn的思想相结合，提出了一个可以广泛应用到各个计算机视觉领域的backbone，在检测、分类和分割等任务的数据集上都呈现出很好的效果，可以应用于很多对精度有较高要求的场景。Swin Transformer之所以能有这么大的影响力主要是因为在 ViT 之后，它通过在一系列视觉任务上的强大表现 ，进一步证明了Transformer是可以在视觉领域取得广泛应用的。

下表中展示了目前swin-t模型在1684X上的性能情况，本文主要针对FP16和INT8模型进行优化部署。

prec time(ms)
FP32 41.890
FP16 7.411
INT8 5.505

2.性能瓶颈分析#

image-20230802174121238

通过bmprofile工具可视化FP16模型在1684X上的运行状态，这里截取了模型中的一个block。从图中可以看出大量的permute(transpose)层穿插其中，一方面带来较大的数据搬运开销，另一方面使得网络无法layergroup，并行效果较差。

3.模型优化#

3.1.transpose消除#

image-20230802174903979

观察图中的attention结构，共有3个transpose层。其中第一个transpose层可以拆解为2个transpose，一是把QKV所在的维度(3)移到了最前面，二是将head所在维度(3)与patch所在维度(49)交换了顺序。由于后面紧跟着split操作是为了将QKV拆分成三个分支，那么此处完全可以不做第一个transpose，而让其直接在原维度上进行split。这样再把第二个transpose的执行改变顺序，让他分别向下移动到三个分支上。这样处理的原因是：在tpu-mlir中是支持transpose与相邻的matmul算子融合的，因此当transpose下移到matmul算子上一层就可以与matmul融合。

细心的读者可能会发现一个问题，QK相乘的matmul其右输入已经有一个transpose了，再叠加另一个transpose在一起还能融合吗？又与哪个transpose融合呢？为了解释这个问题，我们可以从下面这张图来分析。

image-20230802181858484

这是我们预期图优化后达到的效果，可以看到这里matmul是将49x32和32x49这两个矩阵做乘法，64和3可以看作batch。刚好我们的tpu-mlir中是支持hdim_is_batch这种优化的。因此对于这种情况，优化后左右两个transpose都被消除掉，在matmul的输出位置会再新增一个transpose。之后这个matmul再与右面剩下的transpose进行Rtrans融合就可以了。效果如下图所示：

image-20230802182409669

这个输出多出来的transpose可以继续被向下移动至下一个matmul之前，此时网络结构如图：

image-20230802182532966

对于第二个matmul，再次应用hdim_is_batch的优化，消除左右输入的transpose层，之后在输出额外加入的transpose就可以刚好和网络最后的transpose层抵消，至此，所有的transpose都被消除了。

image-20230802182655035

相关代码：tpu-mlir/lib/Dialect/Top/Canonicalize/MatMul.cpp

MatmulWithPermuteAndSplit这个pattern就是用于识别swint中的attention结构，并将transpose+split+squeeze的结构进行调整，其目的就是为了让整块结构可以成功的利用我们编译器已有的一系列针对transpose+matmul这个组合的优化。

3.2.更好的layergroup#

TPU 分为 Local Memory 和 Global Memory，一个独立的计算指令的执行会经过 gdma(global → local)，bdc，gdma(local→ global)的过程，在模型的执行过程中，我们希望gdma搬运类的操作越少越好，这样可以更大程度地利用我们的TPU算力。基于这个想法，在tpu-mlir中设计了LayerGroup的功能，LayerGroup经过计算，可以将多个计算指令划分到一个 Group ，在一个 Group 内，每个 Op 直接使用上一个 Op 计算后存放在 Local Memory的数据 ，可以减少每两个Op数据衔接之间的搬出与搬入，从而减少了 io 的时间。因此，layergroup的效果往往也是我们优化一个模型要考虑的因素。

在完成了3.1中的优化工作后，按照运算逻辑，attention应该可以Group到一起，但实际情况并不如此，如图所示，这里截取了final.mlir的一部分attention结构，这里的各个op都是global layer，说明其中仍存在优化点。

image-20230803105715004

经过在tpu-mlir中的debug，分析其原因有两点，一是SliceOp的local layer不支持5维的情况，二是SqueezeOp没有支持localgen。下面针对这两点进行优化。

3.1.1.SliceOp#

代码：tpu-mlir/lib/Dialect/Tpu/Interfaces/Common/Slice.cpp

其中 LogicalResult tpu::SliceOp::LocalGenSupport()用于判断该Op能否支持locallayer，其中

1
2
3
4
5
6
7
8
9
10
11
12
13
14
else if (module::isBM1684XFamily()) {
    if((int)getRunMode(getOperation()) == 1) {
      return failure();
    }
    const auto offset = module::getI64Array(getOffset());
    const auto steps = module::getI64Array(getSteps());
    if (num_dims > 2) {
      if (steps->at(1) != 1)
        return failure();
    }
    if (num_dims > 4) {
      return failure();
    }
}

在这段代码中观察到，对于1684X处理器，在num_dims>4时直接认为不支持local layer。这里我们对逻辑做进一步完善，在group3d的情况下，5维的shape会按照[n,c,d,h*w]来处理，所以此时如果仅做slice_d，是不会导致数据有跨npu整理的行为的，那么此时他也是允许local layer的。

1
2
3
4
5
6
7
8
9
10
11
12
13
if(num_dims == 5){
    int64_t in_shape[5];
    int64_t out_shape[5];
    tpu_mlir::group_type_t group_type = GROUP_3D;
    module::getNCDHW(getInput(), in_shape[0],in_shape[1],in_shape[2],in_shape[3], in_shape[4],group_type);
    module::getNCDHW(getOutput(), out_shape[0],out_shape[1],out_shape[2],out_shape[3], out_shape[4], group_type);
    for(int i=0; i<5; ++i){
      if(in_shape[i]!=out_shape[i] && (i!=2)){
        return failure();
      }
    }
    return success();
}

3.1.2.SqueezeOp#

SqueezeOp还没有支持local layer的codegen，但是1684X的后端中reshape算子是有local实现的，SqueezeOp刚好可以使用。

首先在TpuOps.td文件中给Tpu_SqueezeOp添加localgen的通用接口定义：DeclareOpInterfaceMethods<LocalGenInterface, [“LocalGenSupport”]>

在这个接口定义的基础上我们需要实现两部分，调用后端算子的接口和判断是否支持local的逻辑。

代码：tpu-mlir/lib/Dialect/Tpu/Interfaces/BM1684X/Squeeze.cpp

这个文件中实现了SqueezeOp调用处理器后端算子的接口，我们为其新增codegen_local_bm1684x的接口。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
void tpu::SqueezeOp::codegen_local_bm1684x(int64_t n_step, int64_t c_step,int64_t h_step, int64_t d_step,int64_t w_step,group_type_t group_type,local_sec_info_t &sec_info) {
  auto op = getOperation();
  auto input_spec = BM168x::get_input_spec(op, group_type);
  auto output_spec = BM168x::get_output_spec(op,group_type);
  if (input_spec->at(0).addr == output_spec->at(0).addr) {
    return;
  }
  auto shape = module::getShape(getOutput());
  reshape_spec_t spec = {0};
  spec.dims = shape.size();
  for (size_t i = 0; i < shape.size(); ++i) {
    spec.shape[i] = shape[i];
  }

  BM168x::call_local_func("backend_api_reshape_local", &spec, sizeof(spec),&sec_info, input_spec->data(), output_spec->data());
}

代码：lib/Dialect/Tpu/Interfaces/Common/Squeeze.cpp

这个文件中实现了SqueezeOp支持localgen的判断逻辑。

1
2
3
4
5
6
7
8
9
10
11
LogicalResult tpu::SqueezeOp::LocalGenSupport() {
  if (module::isCV18xx() || module::isBM1684Family()) {
    return failure();
  }
  auto ishape = module::getShape(getInput());
  auto oshape = module::getShape(getOutput());
  if (ishape.size() < 2 || oshape.size() < 2 || ishape[0] != oshape[0] || ishape[1] != oshape[1]) {
    return failure();
  }
  return success();
}

完成上述优化后让我们再来编译模型看一下效果：

image-20230803113806197

可以看到刚刚几个global layer已经整理到一个group中了。

3.1.3.weight切分#

从上面group的效果来看，还存在着一个比较特殊的情况，这里AddOp并没有和其他层group到一起。这里的原因是，add的一个输入为权重，但是tpu-mlir目前对权重的处理是不进行切分，所以在切分遇到weight时，就认为其不支持group。但是像add这种点对点运算的Op，如果输入为权重，理论上也是可以进行切分的。

为了支持这个功能，涉及修改的地方较多，感兴趣的读者可以先了解一下tpu-mlir中layer group的过程实现，相关的讲解视频在开源社区:layer group精讲

此处概括一下支持weight切分的方式：

1.给top层WeightOp增加 allow_split的参数

2.LocalGenSupport支持add sub mul div max min这类点对点操作

3.做layer group之前给符合要求的op配置allow_split

4.完善layer group切分时涉及到输入为weight的分支逻辑

完成上述优化后让我们再来编译模型看一下效果：

image-20230803142118924

AddOp也成功的合入了group。

4.优化效果#

使用bmprofile工具再次观察模型的运行情况，与优化前相比，节省了大量GDMA搬运的时间，BDC计算与GDMA搬运数据的并行效果更好了。

image-20230803142608913

模型的性能变化情况：

FP16 INT8
优化前 7.411ms 5.505ms
优化后 3.522ms 2.228ms
优化效果 性能提升110% 性能提升145%

FP16模型和INT8模型在1684X上的运行速度都得到了大幅度提升。至此，这一阶段的swint优化工作完成。

希望这篇记录文档能为其他类似模型的优化工作提供帮助。

TPU-MLIR敏感层搜索功能介绍
ChatGLM2-6B解析与TPU部署