阿里雲PAI-Diffusion功能再升級，全鏈路支持模型調優，平均推理速度提升75%以上

語言: CN / TW / HK
時間 2023-02-08 14:01:58

導讀

隨着Stable Diffusion模型在AI圈的爆火，AI生成內容（AI Generated Content，AIGC）的應用呈現出井噴式的增長趨勢。其中，文圖生成（Text-to-image Generation）任務是最流行的AIGC任務之一，旨在生成與給定文本對應的圖像。在先前的工作中，阿里雲機器學習PAI團隊開源了PAI-Diffusion系列模型（看這裏），包括一系列通用場景和特定場景的文圖生成模型，例如古詩配圖、二次元動漫、魔幻現實等。這些模型的Pipeline除了包括標準的Latent Diffusion Model，還集成了PAI團隊先前提出的中文CLIP跨模態對齊模型（看這裏），及圖像超分模型，使得模型可以生成符合中文文本描述的、各種場景下的高清大圖。由於在特定場景下，用户往往有對這些中文Diffusion模型調優（Fine-tune）的需求，而且在線上部署方面，由於Diffusion模型推理速度比較慢，而且需要耗費較多的硬件資源。在本次的工作中，我們在EasyNLP算法框架中，對先前提出的PAI-Diffusion模型功能進行升級，全鏈路支持模型調優。在模型推理方面，結合由PAI自主研發的編譯優化工具 PAI-Blade，支持對PAI-Diffusion模型的端到端的導出和推理加速，在A10機器下做到了1s內的中文大圖生成。

本⽂首先介紹如何體驗PAI-Diffusion模型以及其在線部署、加速推理能力。其次，我們簡單回顧了PAI-Diffusion模型的架構，之後詳細介紹了在EasyNLP算法框架中對上述模型進行調優的全鏈路支持。

輕鬆體驗PAI-Diffusion模型

首先讓我們輕鬆體驗PAI-Diffusion模型。

在線體驗

PAI-Diffusion模型可以用於各種風格的藝術畫的生成。這個模型的參數量僅10億，使用一張P100 GPU顯卡，就可以生成1024*1024的高清大圖，示例如下。

我們在AI模型開源社區ModelScope的創空間上線了我們的一個PAI-Diffusion模型（體驗鏈接），歡迎廣大用户體驗。

PAI-Diffusion模型架構回顧

我們簡單回顧一下PAI-Diffusion模型的技術架構。簡單的説，PAI-Diffusion模型是面向中文的、基於隱式擴散模型（Latent Diffusion Model, LDM）的文圖生成模型。因為中英文在文化和表達上具有很大的差異性，產出的模型通常無法建模中文特有的現象。PAI-Diffusion模型由阿里雲機器學習（PAI）團隊發佈並開源，除了可以用於通用文圖生成場景，還具有一系列特定場景的定製化中文Diffusion模型，包括古詩配圖、二次元動漫、魔幻現實等。PAI-Diffusion的模型Pipeline架構，包括中文CLIP模型、Diffusion模型、圖像超分模型等：

具體地説，中文CLIP模型把中文文本輸入轉化成 Embedding 向量，我們採用EasyNLP中文CLIP跨模態對齊模型（看這裏）的Text Transformer作為Text Encoder。我們把中文CLIP模型的輸出作為LDM的輸入，訓練模型生成文本對應的圖片。可選地，PAI-Diffusion模型鏈路中還集成了ESRGAN作為圖像超分模型，提升圖片的分辨率。

PAI-Diffusion模型加速推理技術

PAI-Blade採用以AI編譯器為主線的技術路徑，支持在多種框架、多種硬件上進行機器學習模型的優化。其核心組件BladeDISC是一款支持動態Shape的深度學習編譯器，並已在GitHub開源（http://github.com/alibaba/BladeDISC）。

image.png

上圖為PAI-Diffusion端到端生成一張圖片的各部分時間佔比，由圖可知，該模型的主要開銷來自UNet中的Resnet Block和Attention Block。針對不同類型的深度學習 workload，BladeDISC有相應的技術進行鍼對性優化。針對計算密集型部分，BladeDISC內置BladNN可以提供高性能的算子實現；針對訪存密集型部分，aStitch則可以對訪存密集型算子進行大尺度的融合；同時，由於BladeDISC原生支持動態Shape，優化完的模型可以輸出不同尺寸的圖片。

通過對PAI-Diffusion模型和社區的Stable-Diffusion模型的推理進行測速分析，PAI-Blade 優化之後成功將整個推理鏈路平均提速75%以上，具體加速比視推理步數略有不同，詳細可以參考下表（測試推理機器 Nvidia-A10，輸出圖片分辨率 512x512）：

pai-diffusion原版 pai-blade 加速後 加速比 採樣步數
1.61s 0.97s 66% 20
3.02s 1.62s 86% 50

綜上可以看到， 引入PAI-Blade 能大幅度提升DiffusionModel的推理速度，降低服務延遲，節省計算資源。

PAI-Diffusion模型的微調和微調後模型部署

在本次更新中，我們在PAI-EasyNLP算法框架對PAI-Diffusion模型的調優做了全面支持。我們介紹如何在PAI-Diffusion模型的訓練流程，以及如何在PAI的產品中部署上述微調後的模型用於在線服務。

安裝EasyNLP

用户可以直接參考鏈接的説明安裝EasyNLP算法框架。

數據準備

首先準備訓練數據與驗證數據，為tsv文件。這⼀⽂件包含以製表符\t分隔的三列（idx, text, imgbase64），第一列是文本編號，第二列是文本，第三列是對應圖片的base64編碼。樣例如下：

64b4109e34a0c3e7310588c00fc9e157	韓國可愛日系襪子女中筒襪春秋薄款純棉學院風街頭卡通兔子長襪潮	iVBORw0KGgoAAAAN...MAAAAASUVORK5CYII=

為了方便開發者，我們也提供了轉換圖片到base64編碼的示例代碼：

import base64
from io import BytesIO
from PIL import Image

img = Image.open(fn)
img_buffer = BytesIO()
img.save(img_buffer, format=img.format)
byte_data = img_buffer.getvalue()
base64_str = base64.b64encode(byte_data) # bytes

下列文件已經完成預處理，可用於測試：

# train
http://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/painter_text2image/T2I_train.tsv

# valid
http://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/painter_text2image/T2I_val.tsv

# test
http://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/painter_text2image/T2I_test.tsv

模型微調

我們採用以下命令對PAI-Diffusion模型進行fine-tune：

easynlp \
    --mode=train \
    --worker_gpu=1 \
    --tables=T2I_train.tsv,T2I_val.tsv \
    --input_schema=idx:str:1,text:str:1,imgbase64:str:1 \
    --first_sequence=text \
    --second_sequence=imgbase64 \
    --checkpoint_dir=./finetuned_model \
    --learning_rate=4e-5 \
    --epoch_num=3 \
    --random_seed=42 \
    --logging_steps=100 \
    --save_checkpoint_steps=1000 \
    --sequence_length=288 \
    --micro_batch_size=16 \
    --app_name=latent_diffusion \
    --user_defined_parameters='
        pretrain_model_name_or_path=alibaba-pai/pai-diffusion-general-large-zh
		reset_model_state_flag=True
      '

訓練完成後模型被保存到./finetune_model/。

模型離線推理

模型訓練完畢後，我們可以將其用於圖片的生成，示例如下：

easynlp \
      --mode=predict \
      --worker_gpu=1 \
      --tables=T2I_test.tsv \
      --input_schema=idx:str:1,text:str:1 \
      --output_schema=text \
      --outputs=./output.tsv \
      --first_sequence=text \
      --checkpoint_dir=./finetuned_model \
      --random_seed=42 \
      --logging_steps=100 \
      --save_checkpoint_steps=500 \
      --sequence_length=32 \
      --micro_batch_size=16 \
      --app_name=latent_diffusion \
      --user_defined_parameters='
      		n_samples=2
			write_image=True
			image_prefix=./output/
    	'

直接在線部署

用户可以使用PAI-EAS自定義服務部署方式，使用如下config部署我們已經打包好的相關processor進行服務部署即可。我們提供了一個測試模型以供體驗模型的推理速度，部署在A10/T4/V100 GPU的config參考如下，您需要將resource字段替換成自己的資源組：

{
    "baseimage": "registry.cn-shanghai.aliyuncs.com/eas/eas-worker-amd64:0.6.8",
    "data_image": "registry.cn-shanghai.aliyuncs.com/pai-ai-test/eas-service:ch_ldm_v100",
    "metadata": {
        "cpu": 15,
        "gpu": 1,
        "instance": 1,
        "memory": 50000,
        "name": "ch_ldm_easynlp",
        "resource": "請替換成自己的資源組",
        "rpc": {
            "keepalive": 50000,
            "worker_threads": 5
        }
    },
    "model_path": "http://pai-vision-exp.oss-cn-zhangjiakou.aliyuncs.com/wzh-zhoulou/dl_eas_processor/ch_ldm/ch_ldm_blade_220206/eas_model_20230206_noblade.tar.gz",
    "processor_entry": "./app.py",
    "processor_path": "http://pai-vision-exp.oss-cn-zhangjiakou.aliyuncs.com/wzh-zhoulou/dl_eas_processor/ch_ldm/ch_ldm_blade_220206/eas_processor_20230206.tar.gz",
    "processor_type": "python",
    "name": "ch_ldm_easynlp"
}

如果您需要體驗Blade優化後的推理速度，可以在A10 GPU中進行部署，用户需要將上述model_path的值替換為http://pai-vision-exp.oss-cn-zhangjiakou.aliyuncs.com/wzh-zhoulou/dl_eas_processor/ch_ldm/ch_ldm_blade_220206/eas_model_20230206.tar.gz。

當服務部署成功後，可以使用如下格式的request body進行測試：

{
    "text": "測試文本",
    "num_inference_steps": 20,
    "num_images": 1
}

在線服務返回示例如下：

{
    "text": "測試文本",
    "images_base64": ["...", "...", ...]
}

images_base64字段返回生成圖像的base64編碼，生成圖像數量由request body中的參數指定。

微調後部署

用户也可以使用EasyNLP進行微調訓練後部署，需要將如下訓練後的ckpt 打包替換上面json中的model_path字段：

打包命令示例如下：

cd finetuned_model/
tar czf finetuned_model.tar.gz config.json pytorch_model.bin RRDB_ESRGAN_x4.pth vocab.txt

然後，採用前述config 將模型部署到PAI-EAS，建議使用的GPU型號為A10/T4/V100。

如果用户需要使用Blade 加速，支持使用A10/A100 GPU機器資源進行部署，請聯繫我們獲取Blade相關教程。

未來展望

在這一期的工作中，我們對PAI-Diffusion模型的功能進行重大升級，在EasyNLP框架中（http://github.com/alibaba/EasyNLP）支持全鏈路模型的調優和在線部署，使得用户能一鍵實現模型的調優，並且在阿里雲PAI平台上便捷部署在線服務。此外，阿里雲機器學習PAI團隊也在持續推進中文多模態模型的自研工作，歡迎用户持續關注我們，也歡迎加入我們的開源社區，共建中文NLP和多模態算法庫！

Github地址：http://github.com/alibaba/EasyNLP

Reference

  1. Chengyu Wang, Minghui Qiu, Taolin Zhang, Tingting Liu, Lei Li, Jianing Wang, Ming Wang, Jun Huang, Wei Lin. EasyNLP: A Comprehensive and Easy-to-use Toolkit for Natural Language Processing. EMNLP 2022
  2. Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer. High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022
  3. Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, Ilya Sutskever. Learning Transferable Visual Models From Natural Language Supervision. ICML 2021

EasyNLP阿里靈傑回顧

EasyNLP是阿里雲機器學習PAI 團隊基於 PyTorch 開發的易用且豐富的中文NLP算法框架，支持常用的中文預訓練模型和大模型落地技術，並且提供了從訓練到部署的一站式 NLP 開發體驗。EasyNLP 提供了簡潔的接口供用户開發 NLP 模型，包括NLP應用 AppZoo 和預訓練 ModelZoo，同時提供技術幫助用户高效的落地超大預訓練模型到業務。由於跨模態理解需求的不斷增加，EasyNLP也將支持各種跨模態模型，特別是中文領域的跨模態模型，推向開源社區，希望能夠服務更多的 NLP 和多模態算法開發者和研究者，也希望和社區一起推動 NLP /多模態技術的發展和模型落地。

「」
「」