利用OpenVINO?部署HuggingFace預訓練模型的方法與技巧

1

背景

作為深度學習領(lǐng)域的 “github” ， HuggingFace 已經(jīng)共享了超過 100,000個預訓練模型，10,000個數(shù)據(jù)集，其中就包括了目前 AIGC 領(lǐng)域非常熱門的“文生圖”，“圖生文”任務范式，例如 ControlNet, StableDiffusion, Blip 等。

通過 HuggingFace 開源的 Transformers, Diffu---sers 庫，只需要要調(diào)用少量接口函數(shù)，入門開發(fā)者也可以非常便捷地微調(diào)和部署自己的大模型任務，你甚至不需要知道什么是 GPT ， BERT 就可以用他的模型，開發(fā)者不需要從頭開始構(gòu)建模型任務，大大簡化了工作流程。從下面的例子中可以看到，在引入 Transformer 庫以后只需要5行代碼就可以構(gòu)建一個基于 GPT2 的問答系統(tǒng)，期間 HuggingFace 會為你自動下載 Tokenizer 詞向量庫與預訓練模型。

圖：HuggingFace 預訓練模型任務調(diào)用示例

但也正因為 Transformer, Diffusers 這些庫具有非常高的易用性，很多底層的代碼與模型任務邏輯也被隱藏了起來，如果開發(fā)者想針對某個硬件平臺做特定的優(yōu)化，則需要將這些庫底層流水行進行拆解再逐個進行模型方面的優(yōu)化。下面這張圖就展示了利用 HuggingFace 庫在調(diào)用 ControlNet 接口時的邏輯, 和他底層實際的流水線結(jié)構(gòu)：

圖：ControlNet 接口調(diào)用邏輯

圖：ControlNet 實際運行邏輯

2

OpenVINO簡介

用于高性能深度學習的英特爾發(fā)行版 OpenVINO 工具套件基于 oneAPI 而開發(fā)，以期在從邊緣到云的各種英特爾平臺上，幫助用戶更快地將更準確的真實世界結(jié)果部署到生產(chǎn)系統(tǒng)中。通過簡化的開發(fā)工作流程， OpenVINO 可賦能開發(fā)者在現(xiàn)實世界中部署高性能應用程序和算法。

在推理后端，得益于 OpenVINO 工具套件提供的“一次編寫，任意部署”的特性，轉(zhuǎn)換后的模型能夠在不同的英特爾硬件平臺上運行，而無需重新構(gòu)建，有效簡化了構(gòu)建與遷移過程。此外，為了支持更多的異構(gòu)加速單元， OpenVINO 的 runtime API 底層采用了插件式的開發(fā)架構(gòu)，基于 oneAPI 中的 oneDNN 等函數(shù)計算加速庫，針對通用指令集進行深度優(yōu)化，為不同的硬件執(zhí)行單元分別實現(xiàn)了一套完整的高性能算子庫，充分提升模型在推理運行時的整體性能表現(xiàn)。

可以說，如果開發(fā)者希望在英特爾平臺上實現(xiàn)最佳的推理性能，并具備多平臺適配和兼容性， OpenVINO是不可或缺的部署工具首選。因此接下來的方案也是在探討如何利用 OpenVINO 來加速 HuggingFace 預訓練模型。

3

OpenVINO 部署方案

簡單來說目前有兩種方案可以實現(xiàn)利用 OpenVINO 加速 Huggingface 模型部署任務，分別是使用 Optimum-Intel 插件以及導出 ONNX 模型部署的方式，兩種方案均有不同的優(yōu)缺點。

圖：OpenVINO 部署 HuggingFace 模型路徑

方案一：使用 Optimum-Intel 推理后端

Optimum-Intel 用于在英特爾平臺上加速 HuggingFace 的端到端流水線。它的 API 和Transformers 或是 Diffusers 的原始 API 極其相似，因此所需代碼改動很小。目前 Optimum-Intel 已經(jīng)集成了 OpenVINO作為其推理任務后端，在大部分 HuggingFace 預訓練模型的部署任務中，開發(fā)者只需要替換少量代碼，就可以實現(xiàn)將HuggingFace Pipeline 中的模型通過 OpenVINO 部署在 Intel CPU 上，并加速推理任務， OpenVINO會自動優(yōu)化 bfloat16 模型，優(yōu)化后的平均延遲下降到了 16.7 秒，相當不錯的 2 倍加速。從下圖可以看到在調(diào)用 OpenVINO的推理后端后，我們可以最大化 Stable Diffusion 系列任務在 Intel CPU 上的推理性能。

圖：Huggingface 不同后端在 CPU 上的性能比較

圖：只需2行代碼替換，利用 OpenVINO 部署文本分類任務

此外 Optimum-Intel 也可以支持在 Intel GPU 上部署模型：

圖：在 Intel GPU 上加載 Huggingface 模型

Optimum Intel 和 OpenVINO 安裝方式如下:

$ pip install optimum[openvino]

在部署 Stable Diffusion 模型任務時，我們也只需要將StableDiffusion Pipeline 替換為 OVStableDiffusionPipeline 即可。

from optimum.intel.openvino import OVStableDiffusionPipeline


ov_pipe = OVStableDiffusionPipeline.from_pretrained(model_id, export=True)

向右滑動查看完整代碼

除此以外 Optimum-Intel 還引入了對 OpenVINO 模型壓縮工具 NNCF 組件的支持，NNCF 目前可以支持 Post-training static quantization （訓練后量化）和 Quantization-aware training （訓練感知量化）兩種模型壓縮模式，前者需要引入少量不帶標簽的樣本數(shù)據(jù)來校準模型輸入的數(shù)據(jù)分布，定制量化參數(shù)，后者則可以在保證模型準確性的情況下，進行量化重訓練。將 HuggingFace 中豐富的數(shù)據(jù)集資源作為校準數(shù)據(jù)或是重訓練數(shù)據(jù)，我們可以輕松完成對預訓練模型的 Int8 在線量化與推理，具體示例如下：

圖：后訓練量化示例

方案二：使用 OpenVINO runtime 進行部署

當然 Optimum-Intel 庫在提供極大便捷性的同時，也有一定的不足，例如對于新模型的支持存在一定的滯后性，并且對 HuggingFace 庫存在依賴性，因此第二種方案就是將 HuggingFace 的預訓練模型直接導出為 ONNX 格式，再直接通過 OpenVINO 的原生推理接口重構(gòu)整個 pipeline ，以此來達到部署代碼輕量化，以及對新模型 pipeline enable 的目的。

這里提供3種導出模型的方案：

使用 Optimum-Intel 接口直接導出OpenVINO 的 IR 格式模型：

圖：使用 Optimum-Intel 直接導出 IR 文件

2. 使用 HuggingFace 原生工具導出 ONNX 格式模型：

HuggingFace 的部分庫中是包含 ONNX 模型導出工具的，以 Transformer 庫為例，我們可以參考其官方文檔實現(xiàn) ONNX 模型的導出。

3. 使用 PyTorch 底層接口導出 ONNX 格式模型：

如果是 Optimum-Intel 還不支持的模型，同時 HuggingFace 庫也沒有提供模型導出工具的話，我們就要通過基礎(chǔ)訓練框架對其進行解析，由于 Transformer 等庫的底層是基于 PyTorch 框架進行構(gòu)建，如何從 PyTorch 框架導出 ONNX 模型的通用方法

這里我們再以 ControlNet 的姿態(tài)任務作為示例，從本文背景章節(jié)中的任務流程圖中我們不難發(fā)現(xiàn) ControlNet 任務是基于多個模型構(gòu)建而成，他的 HuggingFace 測試代碼可以分為以下幾個部分：

1) 加載并構(gòu)建 OpenPose 模型任務

openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')

向右滑動查看完整代碼

2) 運行 OpenPose 推理任務，獲得人體關(guān)鍵點結(jié)構(gòu)

image = openpose(image)

向右滑動查看完整代碼

3) 加載并構(gòu)建 ControlNet 模型任務

controlnet = ControlNetModel.from_pretrained(
  "lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16
)

向右滑動查看完整代碼

4) 下載并構(gòu)建 Stable Diffusion 系列模型任務，并將 ControlNet 對象集成到 StableDiffusion 原始的 pipeline 中

pipe = StableDiffusionControlNetPipeline.from_pretrained(
  "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
)

向右滑動查看完整代碼

5) 運行整個 pipeline 獲取生成的結(jié)果圖像

image = pipe("chef in the kitchen", image, 
num_inference_steps=20).images[0]

向右滑動查看完整代碼

可以看到在1，3，4步任務中夠封裝了模型的下載，因此我們需要對這些接口進行“逆向工程”，找出其中的 PyTorch 的模型對象，并利用 PyTorch 自帶的 ONNX 轉(zhuǎn)換接口 torch.onnx.export(model, (dummy_input, ), 'model.onnx')，將這些對象導出為 ONNX 格式，在這個接口最重要的兩個參數(shù)分別為 torch.nn.Module 模型對象 model ，和一組模擬的輸入數(shù)據(jù) dummy_input，由于 PyTorch 是支持動態(tài)的 input shape ，輸入沒有固定的 shape ，因此我們需要根據(jù)實際情況，找到每個模型的 input shape ，然后再創(chuàng)建模擬輸入數(shù)據(jù)。在這個過程這里我們分別需要找到這個幾個接口所對應庫的源碼，再進行重構(gòu)：

1) OpenPose 模塊

通過解析推理時實際調(diào)用的模型對象，我們可以了解到，這個模型的 PyTorch 對象類為 class bodypose_model(nn.Module)，輸入為 NCHW 格式的圖像 tensor，而他在 controlnet_aux 庫推理過程中抽象出的實例是 OpenposeDetector.body_estimation.model，因此我們可以通過以下方法將他導出為 ONNX 格式：

torch.onnx.export(openpose.body_estimation.model, torch.zeros([1, 3, 184, 136]), OPENPOSE_ONNX_PATH)

向右滑動查看完整代碼

因為 OpenVINO 支持動態(tài)的 input shape，所以 export 函數(shù)中對于輸入的長和寬可以隨機定義。

2) StableDiffusionControlNetPipeline 模塊

圖：StableDiffusionControlNetPipeline 對象初始化參數(shù)

可以看到在構(gòu)建 StableDiffusionControlNetPipeline 的時候，會初始化4個 torch.nn.Module 模型對象，分別是 vae, text_encoder, unet, controlnet, 因此我們在重構(gòu)任務的過程中也需要手動導出這幾個模型對象，此時你必須知道每一個模型的 input shape，以此來構(gòu)建模擬輸入數(shù)據(jù)，這里比較常規(guī)的做法是：直接調(diào)取 pipeline 中的成員函數(shù)進行單個模型的推理任務作為 torch.onnx.export 函數(shù)中的 model 實例。

pipe.text_encoder(
            uncond_input.input_ids.to(device),
            attention_mask=attention_mask,
        )

單獨調(diào)取 text_encoder 推理任務

遍歷 StableDiffusionControlNetPipeline 的__call__函數(shù)，我們也不難發(fā)現(xiàn)，多個模型之間存在串聯(lián)關(guān)系。因此我們也可以模仿 StableDiffusionControlNetPipeline 的調(diào)用任務，構(gòu)建自己的 pipeline，并通過運行這個 pipeline 找到每個模型的 input shape。直白來說就是先重構(gòu)任務，再導出模型：

圖：ControlNet 和 Unet 串聯(lián)

為了更方便地搜索出每個模型的輸入數(shù)據(jù)維度信息，我們也可以為每個模型單獨創(chuàng)建一個“鉤子”腳本，用于替換原始任務中的推理部分的代碼，“鉤取”原始任務的輸入數(shù)據(jù)結(jié)構(gòu)。以 ControlNet 模型為例。

3) 查詢原始模型的輸入?yún)?shù)，將以對應到實際任務的輸入?yún)?shù)。

4) 創(chuàng)建鉤子腳本

class controlnet_input_shape??(object):
  def __init__(self, model) -> None:
    super().__init__()
    self.model = model
    self.dtype = model.dtype
    
  def __call__(self,latent_model_input,
        t,
        encoder_hidden_states,
        controlnet_cond,
        return_dict):
    print("sample:" + str(latent_model_input.shape),
        "timestep:" + str(t.shape),
        "encoder_hidden_states:" + str(encoder_hidden_states.shape),
        "controlnet_cond:" + str(controlnet_cond.shape))
      
  def to(self, device):
    self.model.to(device)

向右滑動查看完整代碼

5) 將鉤子對象替換原來的 controlnet 模型對象，并運行原始的 pipeline 任務

 hooker = controlnet_input_shape--(pipe.controlnet)
pipe.controlnet = hooker

向右滑動查看完整代碼

6)----運行結(jié)果

$ “sample:torch.Size([2, 4, 96, 64]) timestep:torch.Size([]) encoder_hidden_states:torch.Size([2, 77, 768]) controlnet_cond:torch.Size([2, 3, 768, 512])”

向右滑動查看完整代碼

模型導出以及重構(gòu)部分的完整演示代碼可以參考以下示例，這里有一點需要額外注意因為 OpenVINO 的推理接口只支持 numpy 數(shù)據(jù)輸入，而 Diffuers 的示例任務是以 Torch Tensor 進行數(shù)據(jù)傳遞，所以這里建議開發(fā)用 numpy 來重新實現(xiàn)模型的前后處理任務，或是在 OpenVINO 模型輸入和輸入側(cè)提前完成格式轉(zhuǎn)換。

4

總結(jié)

作為當下最火的預訓練模型倉庫之一，HuggingFace可以幫助我們快速實現(xiàn)AIGC 類模型的部署，通過引入Optimum-Intel以及OpenVINO工具套件，開發(fā)者可以更進一步提升這個預訓練模型在英特爾平臺上的任務性能。以下是這兩種方案的優(yōu)缺點比較：

審核編輯：劉清