分割一切？Segment Anything量化加速實戰(zhàn)

01 一、前言

“分割一切,大家一起失業(yè)！”——近期，這樣一句話在社交媒體上大火！這講的就是Segment Anything Model（簡稱 “SAM” ）。SAM 到底是什么？它具備哪些功能？它真的有這么強大嗎？讓我們一起通過本文了解詳情！

SAM 是一個由 Meta AI 實驗室推出的強大人工智能圖像分割應(yīng)用，可以自動識別哪些圖像像素屬于一個對象，并且對圖像中各個對象進行自動風(fēng)格處理，可廣泛用于分析科學(xué)圖像、編輯照片等。

SAM 的完整應(yīng)用由一個圖片編碼器模型（encoder）以及掩碼解碼（mask decoder） + 提示編碼模型（prompt encoder）構(gòu)成，這兩部分都可以被解析為獨立的靜態(tài)模型。其中大部分的算力負載和推理延時都集中在圖片編碼器任務(wù)，因此如果進一步提升圖片編碼器部分的執(zhí)行效率，就成為了 SAM 應(yīng)用的主要優(yōu)化方向之一。

圖：SAM 模型任務(wù)pipeline

本次分享講重點演示如何通過 OpenVINO 的 NNCF 模型壓縮工具實現(xiàn)對 SAM 編碼器部分的量化壓縮，實現(xiàn)在 CPU 側(cè)的性能提升。

02 量化介紹

在正式開始實戰(zhàn)之前，我們不得不提一下量化的概念，量化是指在不改變模型結(jié)構(gòu)的情況下，將模型參數(shù)的表達區(qū)間從 FP32 映射到 INT8 或是 INT4 范圍，用更小數(shù)值位寬來表示相同的信息，實現(xiàn)對于模型體積的壓縮，降低內(nèi)存消耗，同時在模型網(wǎng)絡(luò)的執(zhí)行過程中，系統(tǒng)會自動調(diào)用硬件平臺專門針對低比特數(shù)據(jù)優(yōu)化的指令集或 kernel 函數(shù)，提升性能。

圖：不同精度數(shù)據(jù)的表示位寬

Intel AVX512 VNNI 擴展指令集實現(xiàn)了將原本需要3個時鐘周期才能完成的INT8矩陣點乘與加法運算壓縮到一個時鐘周期，而在最新的 AMX 指令集更是將多個 VNNI 模塊進行堆疊實現(xiàn)了單周期內(nèi)成倍的性能提升。

圖：INT8 矩陣點乘與加法運算指令集優(yōu)化

03 NNCF 訓(xùn)練后量化模式

NNCF 工具的全稱是 Neural Network Compression Framework，是 OpenVINO 工具鏈中專門用于模型壓縮加速的方案實現(xiàn)，包含量化，剪枝，二值化等多種模型壓縮算法，調(diào)用方式又可以分化為訓(xùn)練后量化 (PTQ) 和 訓(xùn)練時壓縮 (QAT) 兩種模式，訓(xùn)練時壓縮要需要引入原始的訓(xùn)練腳本和數(shù)據(jù)集，而訓(xùn)練后量化則可以直接針對訓(xùn)練生成模型文件進行壓縮，無需額外的訓(xùn)練腳本和標注數(shù)據(jù)集參與，這也是 NNCF 在 OpenVINO 2023.0 正式發(fā)布的新功能特性, 而這個模式也僅僅需要以下兩步便可實現(xiàn)：

1.準備校驗數(shù)據(jù)集，這里的校驗數(shù)據(jù)僅用作量化過程中對數(shù)據(jù)表示范圍與分布的計算，因此不需要額外的標簽數(shù)據(jù)，例如在圖像識別任務(wù)中，我們僅需要送入200-300張左右的圖片文件即可。此外我們還需要定義 DataLoader 對象與 transform_fn 數(shù)據(jù)轉(zhuǎn)換函數(shù)， DataLoader 用于讀取校驗數(shù)據(jù)集中的每一個元素，transform_fn 用于將讀取的元素轉(zhuǎn)化為 OpenVINO 模型推理的直接輸入數(shù)據(jù)。

import nncf

calibration_loader = torch.utils.data.DataLoader(...)

def transform_fn(data_item):

images, _ = data_item

return images

calibration_dataset = nncf.Dataset(calibration_loader,transform_fn)

2.運行模型量化，首先需要導(dǎo)入模型對象，然后通過 nncf.quantize() 接口，將模型對象與校驗數(shù)據(jù)集綁定開啟量化任務(wù)， NNCF 工具可以支持多種模型對象類型，包含openvino.runtime.Model, torch.nn.Module, onnx.ModelProto以及 tensorflow.Module

model = ... #OpenVINO/ONNX/PyTorch/TF object

quantized_model = nncf.quantize(model, calibration_dataset)

3.（可選）準確性控制模式，如果發(fā)現(xiàn) NNCF 在默認模式下的導(dǎo)出的模型準確性下降超過預(yù)期，我們也可以使用準確性控制模式（accuracy control）完成訓(xùn)練后量化，此時我們需要加入帶標簽的測試集數(shù)據(jù)，用來評估模型在量化過程中哪些 layer 對模型準確性損失的影響（敏感度）比較大，并作為排序依據(jù)，依次將這些 layer 回退至原始精度，直到模型符合預(yù)期準確性表現(xiàn)。通過這個模式，我們可以在保證模型準確性的情況下，盡可能壓縮模型體積，實現(xiàn)性能和準確性之間的平衡。具體方法可以參考以下鏈接：

https://docs.openvino.ai/nightly/quantization_w_accuracy_control.html

04 Segment Anything + NNCF實戰(zhàn)

接下來讓我們具體一步步看下如何使用 NNCF 的 PTQ 模式完成 SAM encoder 的量化。

1.定義數(shù)據(jù)加載器

本示例使用 coco128 作為校驗數(shù)據(jù)集，其中包含 128 張 .jpg 格式的圖片。由于在量化 ONNX 或 IR 靜態(tài)模型的情況下，數(shù)據(jù)加載器必須是一個 torch 的 DataLoader 類，因此這里我們需要繼承 torch.utils.data.Dataset 并重新構(gòu)建一個數(shù)據(jù)集類，其中必須包含__getitem__方法，用于遍歷數(shù)據(jù)集中的每一個對象，__len__用于獲取數(shù)據(jù)集的對象數(shù)量，最后再通過 torch.utils.data.DataLoader 方法生成數(shù)據(jù)加載器。

classCOCOLoader(data.Dataset):

def__init__(self, images_path):

self.images =list(Path(images_path).iterdir())

def__getitem__(self, index):

image_path =self.images[index]

image =cv2.imread(str(image_path))

image =cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

returnimage

def__len__(self):

returnlen(self.images)

coco_dataset =COCOLoader(OUT_DIR /'coco128/images/train2017')

calibration_loader =torch.utils.data.DataLoader(coco_dataset)

2.定義數(shù)據(jù)格式轉(zhuǎn)化模塊

下一步是定義數(shù)據(jù)轉(zhuǎn)化模塊，我們可以調(diào)用之前定義 preprocess_image 函數(shù)完成數(shù)據(jù)的預(yù)處理，值得注意的是由于 calibration_loader 模塊返回的單個數(shù)據(jù)對象為 torch tensor 類型 ，而 OpenVINO 的 Python 接口不支持該類型數(shù)據(jù)，我們需要先將其強制轉(zhuǎn)化為 numpy 格式。

deftransform_fn(image_data):

image=image_data.numpy()

processed_image=preprocess_image(np.squeeze(image))

returnprocessed_image

calibration_dataset=nncf.Dataset(calibration_loader,transform_fn)

3.運行 NNCF 量化

為了確保量化后的模型準確性，這里我們使用原始的 FP32 ONNX 格式模型作為輸入對象，而不是 FP16 的 IR 格式模型，然后再將該對象送入 nncf.quantize 接口執(zhí)行量化，該函數(shù)接口中有幾個比較重要的額外參數(shù)：

# Load FP32 ONNX model

model=core.read_model(onnx_encoder_path)

quantized_model=nncf.quantize(model,

calibration_dataset,

model_type=nncf.parameters.ModelType.TRANSFORMER,

preset=nncf.common.quantization.structs.QuantizationPreset.MIXED)

ov_encoder_path_int8="sam_image_encoder_int8.xml"

serialize(quantized_model,ov_encoder_path_int8)

model_type：模型類別，用于開啟特殊的量化策略，例如在類 Transformer 模型中，我們需要優(yōu)先保證模型的準確性。

preset：量化模式，默認為 PERFORMANCE，使用對卷積的權(quán)重和偏置均采用對稱量化算法，有助于提升模型性能，此處為了提升模型準確性，我們采用 MIXED 模式，采用權(quán)重對稱量化，偏置非對稱量化的方法，適合模型中包含非 Relu 或者非對稱的激活層。

由于 SAM encoder 模型的網(wǎng)絡(luò)結(jié)構(gòu)比較復(fù)雜，而量化過程中我們需要多次遍歷模型每一個 layer 的參數(shù)，所以量化耗時相對會長一些，請大家耐心等待。這邊建議使用 32G 以上內(nèi)存的硬件設(shè)備，如果遇到內(nèi)存不夠的情況，可以通過 subset_size=100 參數(shù)，適當降低校驗數(shù)據(jù)數(shù)量。

4.模型準確性比較

接下來我們比較下 INT8 和 FP16 模型的推理結(jié)果：

圖：prompt 模式 FP16 – INT8 結(jié)果比較

圖：auto 模式 FP16 – INT8 結(jié)果比較

可以看到在 prompt 和 auto 模式下，INT8 模型的準確性相較 FP16 模型，幾乎沒有任何變化。

注：auto 模式下，mask 將使用隨機生成的顏色。

5.性能比較

最后我們通過 OpenVINO 自帶的 benchmark_app 工具比較下性能指標：

圖：Benchmark 結(jié)果 （FP16）

圖：Benchmark 結(jié)果 （INT8）

可以看到在 CPU 端，INT8 模型相較 FP16 提升了大約 30%， 體積從原本的 350MB 壓縮到了 100MB 不到。

05 總結(jié)

鑒于 SAM 出色的自動化分割能力，相信未來會有越來越多應(yīng)用場景會部署這項技術(shù)，而在產(chǎn)業(yè)化落地的過程中，開發(fā)者往往最關(guān)注的就是性能和準確性之間的平衡，以此獲取成本更優(yōu)的方案。OpenVINO NNCF 工具通過對 Segment Anything encoder 部分的量化壓縮，在幾乎沒有影響模型準確性的情況下，顯著提升模型的運行效率，降低模型占用空間。

審核編輯：劉清