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

導讀

隨着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 採用多種模型編譯優化技術，可以對基於Tensorflow/Pytorch 訓練的模型進行推理優化加速。

 

針對基於PAI-EasyNLP的 PAI-Diffusion模型 和社區的Stable-Diffusion模型的推理進行測速分析，針對推理的主要部分，引入了編譯優化和flashattention的推理加速技術，成功將整個推理鏈路平均提速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
https://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/painter_text2image/T2I_train.tsv

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

# test
https://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框架中（https://github.com/alibaba/EasyNLP）支持全鏈路模型的調優和在線部署，使得用戶能一鍵實現模型的調優，並且在阿里雲PAI平臺上便捷部署在線服務。此外，阿里雲機器學習PAI團隊也在持續推進中文多模態模型的自研工作，歡迎用戶持續關注我們，也歡迎加入我們的開源社區，共建中文NLP和多模態算法庫！

Github地址：https://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
4. Dao T, Fu D Y, Ermon S, et al. Flashattention: Fast and memory-efficient exact attention with io-awareness. arXiv preprint arXiv:2205.14135, 2022

原文鏈接

本文爲阿里雲原創內容，未經允許不得轉載。