SwinTransformer模型優(yōu)化

1.SwinTransformer概述#

自從Transformer在NLP任務(wù)上取得突破性的進(jìn)展之后，業(yè)內(nèi)一直嘗試著把Transformer用于CV領(lǐng)域。之前的若干嘗試都是將Transformer用在了圖像分類領(lǐng)域，但這些方法都面臨兩個(gè)非常嚴(yán)峻的挑戰(zhàn)，一是多尺度問(wèn)題，二是計(jì)算復(fù)雜度的問(wèn)題。

基于這兩個(gè)挑戰(zhàn)，swint的作者提出了一種層級(jí)式提取的Transformer，并通過(guò)移動(dòng)窗口的方式來(lái)學(xué)習(xí)特征。在窗口內(nèi)計(jì)算自注意力可以帶來(lái)更高的效率；同時(shí)通過(guò)移動(dòng)的操作，讓相鄰的窗口之間有了交互，變相達(dá)到了一種全局建模的能力，進(jìn)而解決了上面兩個(gè)問(wèn)題。

Swin Transformer將transformer結(jié)構(gòu)與cnn的思想相結(jié)合，提出了一個(gè)可以廣泛應(yīng)用到各個(gè)計(jì)算機(jī)視覺(jué)領(lǐng)域的backbone，在檢測(cè)、分類和分割等任務(wù)的數(shù)據(jù)集上都呈現(xiàn)出很好的效果，可以應(yīng)用于很多對(duì)精度有較高要求的場(chǎng)景。Swin Transformer之所以能有這么大的影響力主要是因?yàn)樵?ViT 之后，它通過(guò)在一系列視覺(jué)任務(wù)上的強(qiáng)大表現(xiàn) ，進(jìn)一步證明了Transformer是可以在視覺(jué)領(lǐng)域取得廣泛應(yīng)用的。

下表中展示了目前swin-t模型在1684X上的性能情況，本文主要針對(duì)FP16和INT8模型進(jìn)行優(yōu)化部署。

2.性能瓶頸分析#

通過(guò)bmprofile工具可視化FP16模型在1684X上的運(yùn)行狀態(tài)，這里截取了模型中的一個(gè)block。從圖中可以看出大量的permute(transpose)層穿插其中，一方面帶來(lái)較大的數(shù)據(jù)搬運(yùn)開(kāi)銷，另一方面使得網(wǎng)絡(luò)無(wú)法layergroup，并行效果較差。

3.模型優(yōu)化#

3.1.transpose消除#

觀察圖中的attention結(jié)構(gòu)，共有3個(gè)transpose層。其中第一個(gè)transpose層可以拆解為2個(gè)transpose，一是把QKV所在的維度(3)移到了最前面，二是將head所在維度(3)與patch所在維度(49)交換了順序。由于后面緊跟著split操作是為了將QKV拆分成三個(gè)分支，那么此處完全可以不做第一個(gè)transpose，而讓其直接在原維度上進(jìn)行split。這樣再把第二個(gè)transpose的執(zhí)行改變順序，讓他分別向下移動(dòng)到三個(gè)分支上。這樣處理的原因是：在tpu-mlir中是支持transpose與相鄰的matmul算子融合的，因此當(dāng)transpose下移到matmul算子上一層就可以與matmul融合。

細(xì)心的讀者可能會(huì)發(fā)現(xiàn)一個(gè)問(wèn)題，QK相乘的matmul其右輸入已經(jīng)有一個(gè)transpose了，再疊加另一個(gè)transpose在一起還能融合嗎？又與哪個(gè)transpose融合呢？為了解釋這個(gè)問(wèn)題，我們可以從下面這張圖來(lái)分析。

這是我們預(yù)期圖優(yōu)化后達(dá)到的效果，可以看到這里matmul是將49x32和32x49這兩個(gè)矩陣做乘法，64和3可以看作batch。剛好我們的tpu-mlir中是支持hdim_is_batch這種優(yōu)化的。因此對(duì)于這種情況，優(yōu)化后左右兩個(gè)transpose都被消除掉，在matmul的輸出位置會(huì)再新增一個(gè)transpose。之后這個(gè)matmul再與右面剩下的transpose進(jìn)行Rtrans融合就可以了。效果如下圖所示：

這個(gè)輸出多出來(lái)的transpose可以繼續(xù)被向下移動(dòng)至下一個(gè)matmul之前，此時(shí)網(wǎng)絡(luò)結(jié)構(gòu)如圖：

對(duì)于第二個(gè)matmul，再次應(yīng)用hdim_is_batch的優(yōu)化，消除左右輸入的transpose層，之后在輸出額外加入的transpose就可以剛好和網(wǎng)絡(luò)最后的transpose層抵消，至此，所有的transpose都被消除了。

相關(guān)代碼：tpu-mlir/lib/Dialect/Top/Canonicalize/MatMul.cpp

MatmulWithPermuteAndSplit這個(gè)pattern就是用于識(shí)別swint中的attention結(jié)構(gòu)，并將transpose+split+squeeze的結(jié)構(gòu)進(jìn)行調(diào)整，其目的就是為了讓整塊結(jié)構(gòu)可以成功的利用我們編譯器已有的一系列針對(duì)transpose+matmul這個(gè)組合的優(yōu)化。

3.2.更好的layergroup#

TPU 分為 Local Memory 和 Global Memory，一個(gè)獨(dú)立的計(jì)算指令的執(zhí)行會(huì)經(jīng)過(guò) gdma(global → local)，bdc，gdma(local→ global)的過(guò)程，在模型的執(zhí)行過(guò)程中，我們希望gdma搬運(yùn)類的操作越少越好，這樣可以更大程度地利用我們的TPU算力。基于這個(gè)想法，在tpu-mlir中設(shè)計(jì)了LayerGroup的功能，LayerGroup經(jīng)過(guò)計(jì)算，可以將多個(gè)計(jì)算指令劃分到一個(gè) Group ，在一個(gè) Group 內(nèi)，每個(gè) Op 直接使用上一個(gè) Op 計(jì)算后存放在 Local Memory的數(shù)據(jù) ，可以減少每?jī)蓚€(gè)Op數(shù)據(jù)銜接之間的搬出與搬入，從而減少了 io 的時(shí)間。因此，layergroup的效果往往也是我們優(yōu)化一個(gè)模型要考慮的因素。

在完成了3.1中的優(yōu)化工作后，按照運(yùn)算邏輯，attention應(yīng)該可以Group到一起，但實(shí)際情況并不如此，如圖所示，這里截取了final.mlir的一部分attention結(jié)構(gòu)，這里的各個(gè)op都是global layer，說(shuō)明其中仍存在優(yōu)化點(diǎn)。

經(jīng)過(guò)在tpu-mlir中的debug，分析其原因有兩點(diǎn)，一是SliceOp的local layer不支持5維的情況，二是SqueezeOp沒(méi)有支持localgen。下面針對(duì)這兩點(diǎn)進(jìn)行優(yōu)化。

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();
    }
}

在這段代碼中觀察到，對(duì)于1684X芯片，在num_dims>4時(shí)直接認(rèn)為不支持local layer。這里我們對(duì)邏輯做進(jìn)一步完善，在group3d的情況下，5維的shape會(huì)按照[n,c,d,h*w]來(lái)處理，所以此時(shí)如果僅做slice_d，是不會(huì)導(dǎo)致數(shù)據(jù)有跨npu整理的行為的，那么此時(shí)他也是允許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還沒(méi)有支持local layer的codegen，但是1684X的后端中reshape算子是有l(wèi)ocal實(shí)現(xiàn)的，SqueezeOp剛好可以使用。

首先在TpuOps.td文件中給Tpu_SqueezeOp添加localgen的通用接口定義：DeclareOpInterfaceMethods

在這個(gè)接口定義的基礎(chǔ)上我們需要實(shí)現(xiàn)兩部分，調(diào)用后端算子的接口和判斷是否支持local的邏輯。

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

這個(gè)文件中實(shí)現(xiàn)了SqueezeOp調(diào)用芯片后端算子的接口，我們?yōu)槠湫略鯿odegen_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

這個(gè)文件中實(shí)現(xiàn)了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();
}

完成上述優(yōu)化后讓我們?cè)賮?lái)編譯模型看一下效果：

可以看到剛剛幾個(gè)global layer已經(jīng)整理到一個(gè)group中了。

3.1.3.weight切分#

從上面group的效果來(lái)看，還存在著一個(gè)比較特殊的情況，這里AddOp并沒(méi)有和其他層group到一起。這里的原因是，add的一個(gè)輸入為權(quán)重，但是tpu-mlir目前對(duì)權(quán)重的處理是不進(jìn)行切分，所以在切分遇到weight時(shí)，就認(rèn)為其不支持group。但是像add這種點(diǎn)對(duì)點(diǎn)運(yùn)算的Op，如果輸入為權(quán)重，理論上也是可以進(jìn)行切分的。

為了支持這個(gè)功能，涉及修改的地方較多，感興趣的讀者可以先了解一下tpu-mlir中l(wèi)ayer group的過(guò)程實(shí)現(xiàn)，相關(guān)的講解視頻在開(kāi)源社區(qū):layer group精講

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

1.給top層WeightOp增加 allow_split的參數(shù)

2.LocalGenSupport支持add sub mul div max min這類點(diǎn)對(duì)點(diǎn)操作

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

4.完善layer group切分時(shí)涉及到輸入為weight的分支邏輯

完成上述優(yōu)化后讓我們?cè)賮?lái)編譯模型看一下效果：

AddOp也成功的合入了group。

4.優(yōu)化效果#

使用bmprofile工具再次觀察模型的運(yùn)行情況，與優(yōu)化前相比，節(jié)省了大量GDMA搬運(yùn)的時(shí)間，BDC計(jì)算與GDMA搬運(yùn)數(shù)據(jù)的并行效果更好了。

模型的性能變化情況：

FP16模型和INT8模型在1684X上的運(yùn)行速度都得到了大幅度提升。至此，這一階段的swint優(yōu)化工作完成。

希望這篇記錄文檔能為其他類似模型的優(yōu)化工作提供幫助。

審核編輯：湯梓紅