利用OpenVINO?部署HuggingFace預(yù)訓(xùn)練模型的方法與技巧

1

背景

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

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

圖：HuggingFace 預(yù)訓(xùn)練模型任務(wù)調(diào)用示例

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

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

圖：ControlNet 實(shí)際運(yùn)行邏輯

2

OpenVINO簡(jiǎn)介

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

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

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

3

OpenVINO 部署方案

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

圖：OpenVINO 部署 HuggingFace 模型路徑

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

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

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

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

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

圖：在 Intel GPU 上加載 Huggingface 模型

Optimum Intel 和 OpenVINO 安裝方式如下:

$ pip install optimum[openvino]

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

from optimum.intel.openvino import OVStableDiffusionPipeline


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

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

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

圖：后訓(xùn)練量化示例

方案二：使用 OpenVINO runtime 進(jìn)行部署

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

這里提供3種導(dǎo)出模型的方案：

使用 Optimum-Intel 接口直接導(dǎo)出OpenVINO 的 IR 格式模型：

圖：使用 Optimum-Intel 直接導(dǎo)出 IR 文件

2. 使用 HuggingFace 原生工具導(dǎo)出 ONNX 格式模型：

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

3. 使用 PyTorch 底層接口導(dǎo)出 ONNX 格式模型：

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

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

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

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

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

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

image = openpose(image)

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

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

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

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

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

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

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

5) 運(yùn)行整個(gè) pipeline 獲取生成的結(jié)果圖像

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

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

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

1) OpenPose 模塊

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

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

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

因?yàn)?OpenVINO 支持動(dòng)態(tài)的 input shape，所以 export 函數(shù)中對(duì)于輸入的長(zhǎng)和寬可以隨機(jī)定義。

2) StableDiffusionControlNetPipeline 模塊

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

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

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

單獨(dú)調(diào)取 text_encoder 推理任務(wù)

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

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

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

3) 查詢?cè)寄Ｐ偷妮斎雲(yún)?shù)，將以對(duì)應(yīng)到實(shí)際任務(wù)的輸入?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)

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

5) 將鉤子對(duì)象替換原來(lái)的 controlnet 模型對(duì)象，并運(yùn)行原始的 pipeline 任務(wù)

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

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

6)----運(yùn)行結(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])”

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

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

4

總結(jié)

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

審核編輯：劉清