TPU-MLIR中的融合處理

通常，模型的輸入是 預(yù)處理過(guò)的NCHW 格式的張量，但有時(shí)我們可能希望直接將原始圖像輸入模型。 在這種情況下，我們需要將預(yù)處理操作作為模型的一部分。我們目前支持這些類(lèi)型的圖像輸入，它們對(duì)應(yīng)不同的像素和通道格式。 其中一些是專(zhuān)門(mén)針對(duì) CV 系列芯片的，所以這部分我們暫且先不深入探討。

對(duì)于預(yù)處理類(lèi)型，我們目前可以做一些典型的預(yù)處理操作，包括中心裁剪、轉(zhuǎn)置、像素格式變換和歸一化,更多類(lèi)型的操作會(huì)在未來(lái)完成。

那么，我們?cè)鯓硬拍艿玫揭粋€(gè)包含預(yù)處理操作的bmodel呢？

首先，我們像往常一樣進(jìn)行 model_transform 以獲得Top層的mlir 文件。但是在 model_deploy 步驟中，我們需要指定 fuse-preprocess 參數(shù)。 一旦指定了這個(gè)參數(shù)，test_input 應(yīng)該是一個(gè)圖像而不是 numpy 數(shù)組。對(duì)于自定義格式，如果不指定，它會(huì)被默認(rèn)設(shè)置為和preprocesse infor一致。

在我們將它提供給模型之前，我們?nèi)匀恍枰獙?duì)圖像數(shù)據(jù)做一些工作，包括：

調(diào)整圖像到指定尺寸。

將數(shù)據(jù)擴(kuò)展為4維。

并轉(zhuǎn)換為ui8 numpy 數(shù)組。這里數(shù)據(jù)格式是 unsinged int8，因?yàn)橄袼刂捣秶?到255的整數(shù)值。

通過(guò)運(yùn)行帶有 fuse-preprocess 參數(shù)的 model_deploy 接口，TPU-MLIR 將從 inputOp 收集所有必要的預(yù)處理和校準(zhǔn)信息，然后將它們保存在 InputOp 之后的 preprocessOp 中。

如果沒(méi)有提供校準(zhǔn)表，這意味著我們得到一個(gè)不需要任何校準(zhǔn)信息的純浮點(diǎn)型bmodel，在這種情況下我們實(shí)際上也不需要任何校準(zhǔn)信息。

接下來(lái)，我們會(huì)對(duì)InputOp進(jìn)行修改，特別是Type，由于模型的輸入是unsigned INT8張量，所以我們需要將其設(shè)置為ui8的均勻量化類(lèi)型。 您會(huì)發(fā)現(xiàn)量化參數(shù)只是單純?cè)O(shè)置為 1，因?yàn)槲覀儗⒄鎸?shí)的校準(zhǔn)信息保存在preprocessOp中了。接著，func type也要修改，以保證其和inputOp的Type一致。完成所有這些工作后，該模型的輸入現(xiàn)在已修改為圖像。

這里我們便可以開(kāi)始lower到 Tpu Dialect。

在這個(gè)階段，如果我們處于F32等浮點(diǎn)量化模式，在 InputOp 之后會(huì)插入一個(gè) CastOp，將數(shù)據(jù)從整數(shù)轉(zhuǎn)換為浮點(diǎn)數(shù)，以確保類(lèi)型的一致性。由于 inputOp 的量化參數(shù)已經(jīng)被設(shè)置為 1。對(duì)于 f32，我們就相當(dāng)于只是單純進(jìn)行類(lèi)型轉(zhuǎn)換。

接下來(lái)，我們將開(kāi)始插入算子，首先在模型中設(shè)置插入點(diǎn)，然后在輸入圖像為NHWC格式時(shí)插入permuteOp。當(dāng)resize_dim與模型輸入形狀不同時(shí)，SliceOp 會(huì)被用來(lái)進(jìn)行裁剪的工作。

之后我們會(huì)在CastOp后設(shè)置一個(gè)新的插入點(diǎn)，ScaleOp在這里會(huì)被用來(lái)完成歸一化操作。然后，我們還將添加 swapChannelOp 以進(jìn)行 RGB 和 BGR之間的轉(zhuǎn)換。完成所有算子插入工作后，我們將刪除 preprocessOp，它就像一個(gè)占位符一樣，利用完后就該跟它說(shuō)再見(jiàn)了。

但是預(yù)處理融合工作還沒(méi)有完成，這里我們還可以做一些優(yōu)化。

如果模型原本的第一個(gè)算子是Conv2d，那么其filter的輸入通道一定是3，對(duì)應(yīng)輸入圖像數(shù)據(jù)的3個(gè)顏色通道。

所以我們可以通過(guò)簡(jiǎn)單地轉(zhuǎn)換filter的輸入通道來(lái)替換swapChannelOp。至此，預(yù)處理工作結(jié)束。

在INT8 量化模式下，情況就有點(diǎn)不同了，

首先，因?yàn)樗胁僮鞫际莍nt8類(lèi)型，所以不會(huì)有CastOp被插入。

其次，因?yàn)槲覀兿雽?duì)int8數(shù)據(jù)進(jìn)行歸一化操作，并且結(jié)果仍然是 int8，所以我們將插入 scalelutOp 而不是 scaleOp。

ScaleLutOp 就像我們?cè)谟貌檎冶淼姆绞竭M(jìn)行均勻量化，也就是說(shuō)我們預(yù)先計(jì)算出所有 256 個(gè)可能的量化結(jié)果并保存在表中，然后我們可以在推理時(shí)直接從表中得到結(jié)果。
但是由于我們有 3 對(duì)mean和scale，所以我們必須創(chuàng)建3個(gè)表來(lái)對(duì)每個(gè)通道的元素分別進(jìn)行量化。

需要注意的一點(diǎn)是，當(dāng)所有均值為零時(shí)，我們假設(shè)它是unsigned int8量化，否則我們就做signed int8量化。

當(dāng)然，有時(shí)候我們會(huì)遇到所有scales_new都等于1，所有means都為0的情況。 在這種情況下，ScalelutOp 實(shí)際上對(duì)輸入數(shù)據(jù)什么都不做，所以我們將跳過(guò)這個(gè)插入部分。

另外，在TPU-MLIR中我們也可能會(huì)使用mix precision模式，比如在int8量化模式下對(duì)原始的首個(gè)OP做浮點(diǎn)運(yùn)算或者在FP量化模式下對(duì)原始的首個(gè)OP做INT8運(yùn)算。

在前者中，將在 preprocessOp 和原始第一個(gè) op 之間插入一個(gè) CastOp，以將 INT8 數(shù)據(jù)轉(zhuǎn)換為FP，對(duì)于這種情況，fuse preprocess步驟將與 int8 量化模式下相同。

但是對(duì)于后者，我們可以看到兩個(gè)CastOp的作用其實(shí)相互抵消了，所以我們可以將它們刪除并再次像 int8量化模式一樣進(jìn)行fuse preprocess操作。

審核編輯：湯梓紅