作者｜梁德澎

AI 作圖領(lǐng)域的工具一直不盡人意，直到去年 8 月 Stable Diffusion 開(kāi)源，成為AI 圖像生成領(lǐng)域無(wú)可爭(zhēng)辯的劃時(shí)代模型。

(資料圖片)

為了提升其推理效率，OneFlow 首度將 Stable Diffusion 模型加速至“一秒出圖”時(shí)代，極大提升了文生圖的速度，在AIGC領(lǐng)域引發(fā)巨大反響，并得到了 Stability.ai 官方的支持。至今，OneFlow 還在不斷刷新 SOTA 紀(jì)錄。

不過(guò)，由于目前大部分團(tuán)隊(duì)主要是基于翻譯 API + 英文 Stable Diffusion 模型進(jìn)行開(kāi)發(fā)，所以在使用中文獨(dú)特的敘事和表達(dá)時(shí)，英文版模型就很難給出正確匹配的圖片內(nèi)容，這對(duì)部分國(guó)內(nèi)用戶來(lái)說(shuō)不太方便。

為了解決這一問(wèn)題，國(guó)內(nèi)的IDEA 研究院認(rèn)知計(jì)算與自然語(yǔ)言研究中心（IDEA CCNL）也開(kāi)源了第一個(gè)中文版本的“太乙 Stable Diffusion”，基于0.2億篩選過(guò)的中文圖文對(duì)訓(xùn)練。上個(gè)月，太乙 Stable Diffusion 在 HuggingFace 上有近 15 萬(wàn)下載量，是下載量最大的中文 Stable Diffusion。

近期，OneFlow 團(tuán)隊(duì)為太乙 Stable Diffusion 適配了 OneFlow 后端，大大提升了推理性能，也可以做到一秒出圖。不少開(kāi)發(fā)者好奇OneFlow使用了哪些優(yōu)化“秘笈”，后文將進(jìn)行簡(jiǎn)要解讀。

歡迎Star、運(yùn)行 OneFlow 版太乙 Stable Diffusion： ?

https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion#without-docker

1對(duì)比 PyTorch，OneFlow 將“太乙 Stable Diffusion”推理速度提升1倍以上

下面的圖表分別展示了在 A100 (PCIe 40GB / SXM 80GB)，V100 ( SXM2 32GB )， RTX 2080，RTX 3080 Ti，RTX 3090， 和 T4 不同類型的 GPU 硬件上分別使用 PyTorch, 和 OneFlow對(duì) 太乙 Stable Diffusion 進(jìn)行推理的性能表現(xiàn)。

可以看到，對(duì)于 A100 顯卡，無(wú)論是 PCIe 40GB 的配置還是 SXM 80GB 的配置，OneFlow 的性能相比 PyTorch 能提升 1 倍以上，推理速度達(dá)到了 50it/s 以上，生成一張圖片所需要的時(shí)間在 1 秒以內(nèi)。

其他硬件數(shù)據(jù)：

注：3090上的AIT數(shù)據(jù)由 IDEA 研究院提供

綜上，在各種硬件的對(duì)比中，對(duì)比 PyTorch， OneFlow 能將太乙 Stable Diffusion 的推理性能提升 1 倍多。

2

生成圖片展示

滔滔江水, 連綿不絕, 唯美, 插畫(huà)

長(zhǎng)城, 清晨, 朦朧, 唯美, 插畫(huà)

夢(mèng)回江南，中國(guó)古代小鎮(zhèn)，唯美，插畫(huà)

中國(guó)的未來(lái)城市, 科幻插畫(huà)

古代建筑, 白雪紛飛

螺螄粉

注：上述圖片均基于 OneFlow 版太乙 Stable Diffusion 生成

3

無(wú)縫兼容 PyTorch 生態(tài)

想體驗(yàn) OneFlow 版的太乙 Stable Diffusion？只需要修改兩行代碼：

之所以能這么輕松遷移模型，是因?yàn)?OneFlow Stable Diffusion 有兩個(gè)出色的特性：

OneFlowStableDiffusionPipeline.from_pretrained 能夠直接使用 PyTorch 權(quán)重。

OneFlow 本身的 API 和 PyTorch 對(duì)齊，因此 import oneflow as torch 之后，torch.autocast、torch.float16 等表達(dá)式完全不需要修改。

上述特性使得 OneFlow 兼容了 PyTorch 的生態(tài)，這不僅在 OneFlow 對(duì) 太乙 Stable Diffusion 的遷移中發(fā)揮了作用，也大大加速了 OneFlow 用戶遷移其它許多模型，比如在和 torchvision 對(duì)標(biāo)的 flowvision 中，許多模型只需通過(guò)在 torchvision 模型文件中加入 import oneflow as torch 即可得到。

此外，OneFlow 還提供全局 “mock torch” 功能，在命令行運(yùn)行 eval $(oneflow-mock-torch) 就可以讓接下來(lái)運(yùn)行的所有 Python 腳本里的 import torch 都自動(dòng)指向 oneflow。 ?

4

動(dòng)靜一體的編程體驗(yàn)

深度學(xué)習(xí)算法原型開(kāi)發(fā)階段需要快速修改和調(diào)試，動(dòng)態(tài)圖執(zhí)行（Eager mode, define by run）最優(yōu)。但在部署階段，模型已經(jīng)固定下來(lái)，計(jì)算效率變得更重要，靜態(tài)圖執(zhí)行（Lazy mode，define and run）可以借助編譯器做靜態(tài)優(yōu)化來(lái)獲得更好的性能。因此，推理階段主要使用靜態(tài)圖模式。

最近，PyTorch 升級(jí)到2.0引入了compile()這個(gè)API，可以把一個(gè)模型或一個(gè)Module從動(dòng)態(tài)圖執(zhí)行變成靜態(tài)圖執(zhí)行。OneFlow里也有一個(gè)類似的機(jī)制，不過(guò)接口名是nn.Graph()，它可以把傳入Module轉(zhuǎn)成靜態(tài)圖執(zhí)行模式。

不僅如此，OneFlow的nn.Graph模式基于MLIR實(shí)現(xiàn)了一系列計(jì)算圖的圖層優(yōu)化，譬如內(nèi)存布局、算子融合等。 ?

這不僅使得計(jì)算圖表示的深度學(xué)習(xí)模型可以在各種硬件上達(dá)到最高性能，更重要的是，使得深度學(xué)習(xí)框架導(dǎo)入的計(jì)算圖更方便地在不同硬件之間實(shí)現(xiàn)遷移，有助于克服國(guó)產(chǎn)硬件軟件生態(tài)薄弱的問(wèn)題。未來(lái)，我們將發(fā)布更多內(nèi)容來(lái)揭示OneFlow深度學(xué)習(xí)編譯器的設(shè)計(jì)和實(shí)現(xiàn)。

歡迎Star、運(yùn)行 OneFlow 版太乙 Stable Diffusion：

https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion#without-docker ?

OneFlow 地址：https://github.com/Oneflow-Inc/oneflow/ ?

其他人都在看

35張圖，直觀理解Stable Diffusion

ChatGPT的一小步，NLP范式轉(zhuǎn)變的一大步

2023年AI十大展望：GPT-4領(lǐng)銜大模型變革

李白：你的模型權(quán)重很不錯(cuò)，可惜被我沒(méi)收了

OpenAI掌門(mén)Sam Altman：AI下一個(gè)發(fā)展階段

比快更快，開(kāi)源Stable Diffusion刷新作圖速度

OneEmbedding:單卡訓(xùn)練TB級(jí)推薦模型不是夢(mèng)

歡迎Star、試用OneFlow最新版本：https://github.com/Oneflow-Inc/oneflow/https://github.com/Oneflow-Inc/oneflow/

關(guān)鍵詞：