Segment Anything量化加速有多強(qiáng)！

前言

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

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

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

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

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

01

量化介紹

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

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

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

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

02

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

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

1. 準(zhǔn)備校驗(yàn)數(shù)據(jù)集

這里的校驗(yàn)數(shù)據(jù)僅用作量化過(guò)程中對(duì)數(shù)據(jù)表示范圍與分布的計(jì)算，因此不需要額外的標(biāo)簽數(shù)據(jù)，例如在圖像識(shí)別任務(wù)中，我們僅需要送入 200-300 張左右的圖片文件即可。此外我們還需要定義 DataLoader 對(duì)象與 transform_fn 數(shù)據(jù)轉(zhuǎn)換函數(shù)， DataLoader 用于讀取校驗(yàn)數(shù)據(jù)集中的每一個(gè)元素，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)

向右滑動(dòng)查看完整代碼

2. 運(yùn)行模型量化

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

model = ... #OpenVINO/ONNX/PyTorch/TF object
quantized_model = nncf.quantize(model, calibration_dataset) 

向右滑動(dòng)查看完整代碼

3.（可選）準(zhǔn)確性控制模式

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

04

Segment Anything + NNCF 實(shí)戰(zhàn)

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

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

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

class COCOLoader(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)
    return image
  
  def __len__(self):
    return len(self.images)
  
coco_dataset = COCOLoader(OUT_DIR / 'coco128/images/train2017')
calibration_loader = torch.utils.data.DataLoader(coco_dataset)

向右滑動(dòng)查看完整代碼

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

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

def transform_fn(image_data):
  image = image_data.numpy()
  processed_image = preprocess_image(np.squeeze(image))
  return processed_image


calibration_dataset = nncf.Dataset(calibration_loader, transform_fn)

向右滑動(dòng)查看完整代碼

3. 運(yùn)行 NNCF 量化

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

向右滑動(dòng)查看完整代碼

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

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

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

4. 模型準(zhǔn)確性比較

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

左右滑動(dòng)查看 prompt 模式 FP16 – INT8 結(jié)果比較

左右滑動(dòng)查看 auto 模式 FP16 – INT8 結(jié)果比較

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

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

5. 性能比較

最后我們通過(guò)OpenVINO自帶的 benchmark_app 工具比較下性能指標(biāo)：

[ INFO ] Execution Devices:['CPU']
[ INFO ] Count:      60 iterations
[ INFO ] Duration:     75716.93 ms
[ INFO ] Latency:
[ INFO ]  Median:    14832.33 ms
[ INFO ]  Average:    14780.77 ms
[ INFO ]  Min:      10398.47 ms
[ INFO ]  Max:      16725.65 ms
[ INFO ] Throughput:  0.79 FPS

Benchmark結(jié)果（FP16）

[ INFO ] Execution Devices:['CPU']
[ INFO ] Count:      72 iterations
[ INFO ] Duration:     68936.14 ms
[ INFO ] Latency:
[ INFO ]  Median:    11281.87 ms
[ INFO ]  Average:    11162.87 ms
[ INFO ]  Min:      6736.09 ms
[ INFO ]  Max:      12547.48 ms
[ INFO ] Throughput:  1.04 FPS

Benchmark 結(jié)果 （INT8）

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

05

總結(jié)

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

審核編輯：劉清