大象轉(zhuǎn)身，TPU-MLIR適配DragGAN模型前向操作

DragGAN

DragGAN是由谷歌、麻省理工學(xué)院和馬克斯普朗克研究所創(chuàng)建的一種新的人工智能模型。

通過點(diǎn)擊、拖動(dòng)等簡(jiǎn)單的交互操作就能改變拍攝對(duì)象的姿勢(shì)、形狀和表情等。

DragGAN改變了傳統(tǒng)的PS操作流程，只需簡(jiǎn)單拖拽起點(diǎn)和終點(diǎn)，AI就會(huì)根據(jù)圖像的變化自動(dòng)生成和補(bǔ)全圖像。

DragGAN可處理的圖像類型豐富多樣，無論是人類表情的調(diào)整還是自然風(fēng)景的變化，都可以在瞬息之內(nèi)輕松實(shí)現(xiàn)。

DragGAN的全流程包含一個(gè)基于Generator的前向操作和反向傳播過程。本文主要介紹在TPU-MLIR上適配DragGAN模型的前向操作的全部過程。

模型移植

推理代碼定位與模型導(dǎo)出

適配的模型代碼使用 XingangPan/DragGAN: Official Code for DragGAN (SIGGRAPH 2023) (github.com) ，模型的入口在 DragGAN/viz/renderer.py:357，可以在這里直接引入TPU-MLIR提供的 gen_shell 工具，直接 trace 生成 workspace 文件夾，onnx/pt 模型，以及默認(rèn)的轉(zhuǎn)換腳本：

fromutils.gen_shellimportgenerate
generate(
"DragGan",
G,
dict(
ws=ws,
c=label
),
"../draggan_workspace",
)

運(yùn)行源碼 README.md 中提供的腳本 python visualizer_drag_gradio.py，運(yùn)行成功后可以在在同級(jí)目錄下得到如下的目錄結(jié)構(gòu)：

draggan_workspace
├──cali_data
│└──data.npz
├──convert.sh
├──DragGan.onnx
├──DragGan.pt
├──data.npz
└──cali_data

模型移植過程中錯(cuò)誤的分析和解決

RuntimeError: Op not support:{'RandomNormalLike'}

在 model_transform 階段，發(fā)現(xiàn)存在不支持的算子 RandomNormalLike：

RandomNormalLike（隨機(jī)數(shù)相關(guān)）的算子 1684x 無法支持，所以必須嘗試在原模型中避開這些算子。定位到模型代碼處，發(fā)現(xiàn)該算子用于提供一個(gè)噪音供下游使用。源碼中提供了三種噪音生成方式，分別是 random（隨機(jī)噪音），const（常量噪音），和 none（不提供噪音），因此可以通過設(shè)置 noise_mode = const 避開這一算子的使用。

對(duì) Conv/DeConv Filter 為動(dòng)態(tài)輸入情況的支持

DragGan 的模型結(jié)構(gòu)中，有一部分 Conv 和 DeConv 的輸入是固定權(quán)重，而 FilterOp 部份是動(dòng)態(tài)的從上游計(jì)算得到的輸入。這種情況在這之前未做考慮，需要添加支持。這包括在多個(gè)地方的代碼更改。下面通過具體的報(bào)錯(cuò)提示來一步步分析、定位和解決。

model_transform 階段

在 tpu-mlir 的 Converter 中，權(quán)重（weight）和 動(dòng)態(tài)輸入（dynamic input）存儲(chǔ)在不同的變量中，其中，weight 通過 getWeightOp(name) 獲取，input 通過 getOperand(name) 獲取。如果不確定 op 是 dynamic input 還是 weight，可以使用 getOp(name) 來獲取。而在對(duì) DragGan 的 model_transform.py 腳本的運(yùn)行過程中，會(huì)遇到如下的報(bào)錯(cuò)KeyError: '/synthesis/b8/conv0/Transpose_output_0'

此時(shí)對(duì)應(yīng)模型結(jié)構(gòu)，發(fā)現(xiàn)該 DeConv 的輸入 /synthesis/b8/conv0/Transpose_output_0 是作為一個(gè) Weight 獲取的。

因此將ConvTranspose 的 filter_opd 的獲取邏輯改為 getOp 即可

同理，另外一個(gè) KeyError 中，DeConv 的 filter 來自于動(dòng)態(tài)輸入，所以同理，將 DeConv 獲取 filter 結(jié)點(diǎn)的邏輯同樣改為 getOp。

在 model_transform 階段， 模型會(huì)首先轉(zhuǎn)換到DragGAN_origin.mlir，再經(jīng)過--shape-infer，--canonicalize 等過程，轉(zhuǎn)換為可以通過model_runner.py做推理的 Top Dialect 描述的 mlir 文件。在對(duì) Top 層做推理驗(yàn)證正確性時(shí)，DragGan 模型報(bào)出了精度為零的錯(cuò)誤。通過觀察輸出的錯(cuò)誤信息，發(fā)現(xiàn)是在 DeConv 層之后精度出現(xiàn)問題的，而且僅在 DeConv 的 filter 是動(dòng)態(tài)輸入的情況下會(huì)有這一問題。

構(gòu)建了一個(gè) filter 是動(dòng)態(tài)輸入的 DeConv 作為單側(cè)，復(fù)現(xiàn)該錯(cuò)誤成功：

classDeConvCase(nn.Module):
def__init__(self)->None:
super().__init__()
self.deconv=nn.ConvTranspose2d(4,4,[2,2],stride=[1,1],bias=False)
self.deconv.weight.data=weight

defforward(self,x,y):
output_padding=self.deconv._output_padding(
x,
None,
[2,2],
[0,0],
[2,2],
1,
[1,1],
)

out=F.conv_transpose2d(x,y,None,[1,1],0,output_padding,1,1)

returnout,self.deconv(x)

此時(shí)通過斷點(diǎn)調(diào)試，發(fā)現(xiàn)錯(cuò)誤原因有兩個(gè)：

• 正確性驗(yàn)證階段推理時(shí)，在 init() 時(shí)設(shè)置權(quán)重，此時(shí) weight 還沒有設(shè)置
• 動(dòng)態(tài)輸入時(shí)沒有做對(duì)應(yīng)的權(quán)重重排（WeightReorder）

tpu-mlir 在適配模型的過程會(huì)經(jīng)過多步轉(zhuǎn)換和多次優(yōu)化，為了保證轉(zhuǎn)換后的正確性，tpu-mlir 會(huì)做三次正確性驗(yàn)證，分別針對(duì) Top Dialect，Tpu Dialect 和 bmodel。Top 和 Tpu 層的正確性的核心代碼位于 ModuleInterpreter.[h/cpp]，該過程會(huì)從輸入開始，對(duì)每一個(gè) Op 分配空間，進(jìn)行初始化（init），在初始化結(jié)束后進(jìn)行推理（inference），并在最終對(duì)每個(gè) Op 進(jìn)行析構(gòu)（deinit）。而 DeConv 的精度錯(cuò)誤之一則來自于 Inference 階段時(shí) init 和 inference 的分離。

在 init 時(shí)，DeConv 會(huì)構(gòu)造一個(gè) Dnnl 的實(shí)例，此時(shí)會(huì)直接 copy 一份 Weight 在 Dnnl 實(shí)例中，但由于該 filter 為動(dòng)態(tài)輸入， init 時(shí)值還沒有傳入，所以傳入的 filter 的值實(shí)質(zhì)上是全零。導(dǎo)致在 inference 階段出現(xiàn)錯(cuò)誤。定位后該問題比較好改，將 init 過程中對(duì) Dnnl 實(shí)例的 setup 移到 inference 階段即可。Conv 也有同樣的問題，修改邏輯相同。

對(duì) onnx 模型，DeConv 的 filter 的權(quán)重存儲(chǔ)方式是 input channel first（即 shape 為 [ic, oc, kw, kh]），而后端的計(jì)算過程大多都需要 output channel first（[oc, ic, kw, kh]），可以注意到 OnnxConverter 中，原本對(duì) DeConv 的權(quán)重會(huì)存在一個(gè)轉(zhuǎn)置操作：

而動(dòng)態(tài)權(quán)重自然沒有辦法實(shí)現(xiàn)這一操作。因此，需要添加一個(gè)圖優(yōu)化，當(dāng) DeConv 的 filter 是動(dòng)態(tài)時(shí)，在其前面添加一個(gè) [oc, ic] 互換的 Permute 操作。在添加 Permute 操作時(shí)，需要仔細(xì)考慮 DeConv 添加這一 Permute 的先決條件。確保該 Permute 添加是針對(duì) DeConv 的動(dòng)態(tài)權(quán)重，且同時(shí)不會(huì)重復(fù)添加。因此考慮在 DeConv 的 Operation 結(jié)構(gòu)中添加 bool 類型的 dynweight_reorderd 參數(shù)。當(dāng) filter 不是 top.WeightOp （使用動(dòng)態(tài)權(quán)重）且 dynweight_reordered 為 false （沒有添加對(duì)動(dòng)態(tài) weight 的 Permute）時(shí)，添加這一 Permute，同時(shí)設(shè)置 dynweight_reorderd 參數(shù)為 true。

在 TopOps.td 文件對(duì) DeConv 添加 dynweight_reorderd 參數(shù)后，對(duì) DeConv 動(dòng)態(tài)權(quán)重的圖優(yōu)化邏輯如下：

structReorderDynWeight:publicOpRewritePattern{
usingOpRewritePattern::OpRewritePattern;

LogicalResultmatchAndRewrite(DeconvOpop,
PatternRewriter&rewriter)constoverride{

autofilter_shape=module::getShape(op.getFilter());//or

if(module::isWeight(op.getOperand(1))){
returnfailure();
}
booldyn_weight_reorderd=op.getDynweightReorderd();
if(dyn_weight_reorderd){
returnfailure();
}

if(isa(op.getOperand(1).getDefiningOp())){
autopermute_op=
dyn_cast(op.getOperand(1).getDefiningOp());

//eraseifalreadyhavethispermutebutfromoriginalgraph
std::vector<int64_t>ps={1,0,2,3};
autoorder=module::getI64Array(permute_op.getOrder());
if(*order==ps){
permute_op.replaceAllUsesWith(permute_op.getInput());
rewriter.eraseOp(permute_op);
op.setDynweightReorderd(true);
returnsuccess();
}
}

rewriter.setInsertionPointAfterValue(op.getFilter());
std::stringname=module::getName(op.getOutput()).str();
autoloc=
NameLoc::get(rewriter.getStringAttr(name+"_reorder_permute"));

std::vector<int64_t>order={1,0};
autofilter_dim=filter_shape.size();
for(inti=2;iorder.push_back(i);
}

autop_type=
UnrankedTensorType::get(module::getElementType(op.getFilter()));
std::vectorattrs;
attrs.emplace_back(
rewriter.getNamedAttr("order",rewriter.getI64ArrayAttr(order)));

autonew_permute_op=rewriter.create(
loc,p_type,ValueRange{op.getFilter()},attrs);

new_permute_op.shape_inference();
op.setOperand(1,new_permute_op.getOutput());
op.setDynweightReorderd(true);
returnsuccess();
}
};

這里做了一個(gè)額外的判斷，當(dāng) DeConv 的 filter 位置已經(jīng)是 Permute 且其 order 和要添加的 Permute 一樣（1,0,2,3）時(shí)，兩個(gè) Permute 可以直接融合，所以此時(shí)可以直接刪除該 Permute 并返回。其他的情況則是插入一個(gè)額外的 Permute 操作。Conv 層同樣要支持動(dòng)態(tài) weight 的權(quán)重重排，要添加一個(gè)相同的圖優(yōu)化。

此外，Top 層的 shape-infer 要早于圖優(yōu)化，因此在做 shape-infer 時(shí)動(dòng)態(tài) weight 的 shape 仍然還是 input channle first，所以 DeConv 的 output_shape 的 dim[1] 應(yīng)該基于 filter_shape[1] 來判斷。對(duì)應(yīng)的修改位于 lib/Dialect/Top/Interfaces/Deconv.cpp：

bmodel 運(yùn)行錯(cuò)誤 ASSERT /workspace/nntoolchain/TPU1686/bm1684x/cmodel/src/cmodel_common.cpp: gather_data: 207: dst_offset < (1<<18)

在大模型中定位這一錯(cuò)誤較難，因此可以通過 mlir_cut.py 逐步縮小范圍，得到了最小可復(fù)現(xiàn)的 mlir：

mlir_cut.py--mlir*tpu.mlir--output_names/synthesis/b64/conv0/Conv_output_0_Conv--input_names/synthesis/b32/conv1/Mul_3_output_0_Mul,/synthesis/b64/conv0/Reshape_3_output_0_Reshape

tpuc-optDragGan_bm1684x_f32_final.mlir--codegen="model_file=DragGan_f32.bmodelembed_debug_info=true"-o/dev/null
model_runner.py--inputfake_data.npz--modelDragGan_f32.bmodel--outputDragGan_bm1684x_f32_model_outputs.npz

進(jìn)一步構(gòu)建了能夠復(fù)現(xiàn)該錯(cuò)誤的單元測(cè)試：

通過控制變量，得到了以下現(xiàn)象：

• 關(guān)閉 layer-group，模型運(yùn)行正常不報(bào)錯(cuò)：這說明問題基本是出在 tpu-mlir 部份而不是后端算子部份
• 將上述的代碼中 DeConv 的 filter 從動(dòng)態(tài)改為靜態(tài)，模型運(yùn)行正常：說明問題仍然是動(dòng)態(tài) Weight 導(dǎo)致的
• 構(gòu)建基本的 DeConv 算子，無論是靜態(tài)和動(dòng)態(tài)都運(yùn)行正常，和上面的單側(cè)進(jìn)行對(duì)比，發(fā)現(xiàn)區(qū)別在單個(gè) DeConv 算子不會(huì)進(jìn)行 LayerGroup：將問題定位到 tpu-mlir 的 LayerGroup 部份的代碼

此時(shí)進(jìn)一步對(duì)比正常和出錯(cuò)的 final.mlir，發(fā)現(xiàn) dynamic weight 和 weight 的 slice 屬性不一致，如下所示：

top.Weight 的 layer-group 是比較特殊。top.Weight 在整個(gè) layer-group 都保存在 local memory 中（hold_in_lmem = true）；同時(shí)，weight 也不能切分 slice，每個(gè) slice 都要用到完成的 filter，從而導(dǎo)致結(jié)果錯(cuò)誤。

所以需要單獨(dú)針對(duì) dynamic weight 處理，這包括設(shè)置其生命周期（hold_in_mem = true），以及將其 slice 設(shè)置為長(zhǎng)度為 1，元素為其 shape 對(duì)應(yīng)維度值的列表。這一過程可以在 lib/Dialect/Tpu/Transforms/LayerGroup/LayerGroupUtil.cpp 的 backward_update_slice 方法中完成：

優(yōu)化后再對(duì)比兩個(gè)單例的 final.mlir，發(fā)現(xiàn)此時(shí) dynamic weight 的 slice 信息已經(jīng)和普通 weight 完全相同：

F16 和 int8 精度問題

在解決了 F32 的 bug 后，F(xiàn)16 和 int8 的 tpu 層 mlir 仍然存在精度問題。原本以為是 DeConv 的 F16 適配存在問題，通過使用 mlir_debugger 對(duì)每一層用正確數(shù)值做推理（也可以直接觀察輸出的 npz 文件以及 npz_tool 的比對(duì)結(jié)果），發(fā)現(xiàn)出錯(cuò)的是 Active -> Mul 的結(jié)構(gòu)，Active 是 ReduceSum 操作：

因此基本可以確定是普通的 F16 溢出問題。驗(yàn)證 BF16，發(fā)現(xiàn)BF16 編譯成功，進(jìn)一步確認(rèn)是溢出問題。

將這些層添加到 qtable 中，發(fā)現(xiàn)還是通過不了比對(duì)，值里面仍然會(huì)存在 inf。對(duì)比發(fā)現(xiàn)是在 Active(ReduceSum) -> Sqrt 的結(jié)構(gòu)中間有兩個(gè) Cast 導(dǎo)致的：

這兩個(gè) cast 沒有作用，可以被優(yōu)化掉，于是寫圖優(yōu)化將這兩個(gè) Cast 直接消除。優(yōu)化后對(duì)應(yīng)的 mlir 如下：

int8 也是相同的精度溢出問題，同樣確認(rèn)兩個(gè) cast 融合的操作能夠覆蓋 int8 的情況即可。

至此，DragGan 適配的模型部份適配完成。

總結(jié)

• 在一些情況下，在不影響結(jié)果的情況下直接修改模型的代碼結(jié)構(gòu)可以更容易的解決一些算子適配問題
• 較大的模型測(cè)試錯(cuò)誤定位到具體算子的情況下，優(yōu)先考慮構(gòu)建單側(cè)嘗試復(fù)現(xiàn)問題
• 控制變量，設(shè)置對(duì)照，是缺少解決思路時(shí)尋找問題的一個(gè)較為通用的方案。