关于人工智能:零代码改动静态编译让太乙Stable-Diffusion推理速度翻倍

4次阅读

共计 2263 个字符，预计需要花费 6 分钟才能阅读完成。

作者|梁德澎

AI 作图畛域的工具始终不尽如人意，直到去年 8 月 Stable Diffusion 开源，成为 AI 图像生成畛域无可争辩的划时代模型。

为了晋升其推理效率，OneFlow 首度将 Stable Diffusion 模型减速至“一秒出图”时代，极大晋升了文生图的速度，在 AIGC 畛域引发微小反应，并失去了 Stability.ai 官网的反对。 至今，OneFlow 还在一直刷新 SOTA 纪录。

不过，因为目前大部分团队次要是基于翻译 API + 英文 Stable Diffusion 模型进行开发，所以在应用中文独特的叙事和表白时，英文版模型就很难给出正确匹配的图片内容，这对局部国内用户来说不太不便。

为了解决这一问题，国内的 IDEA 研究院认知计算与自然语言钻研核心（IDEA CCNL）也开源了第一个中文版本的“太乙 Stable Diffusion”，基于 0.2 亿筛选过的中文图文对训练。 上个月，太乙 Stable Diffusion 在 HuggingFace 上有近 15 万下载量，是下载量最大的中文 Stable Diffusion。

近期，OneFlow 团队为太乙 Stable Diffusion 适配了 OneFlow 后端，大大晋升了推理性能，也能够做到一秒出图 。不少开发者好奇 OneFlow 应用了哪些优化“秘笈”，后文将进行简要解读。

欢送 Star、运行 OneFlow 版太乙 Stable Diffusion：

https://github.com/Oneflow-In…

上面的图表别离展现了在 A100 (PCIe 40GB / SXM 80GB)，V100 (SXM2 32GB)，RTX 2080，RTX 3080 Ti，RTX 3090，和 T4 不同类型的 GPU 硬件上别离应用 PyTorch, 和 OneFlow 对太乙 Stable Diffusion 进行推理的性能体现。

能够看到，对于 A100 显卡，无论是 PCIe 40GB 的配置还是 SXM 80GB 的配置，OneFlow 的性能相比 PyTorch 能晋升 1 倍以上，推理速度达到了 50it/s 以上，生成一张图片所须要的工夫在 1 秒以内。

其余硬件数据：

注：3090 上的 AIT 数据由 IDEA 研究院提供

综上，在各种硬件的比照中，比照 PyTorch，OneFlow 能将太乙 Stable Diffusion 的推理性能晋升 1 倍多。

滔滔江水, 连绵不绝, 唯美, 插画

长城, 凌晨, 朦胧, 唯美, 插画

梦回江南，中国现代小镇，唯美，插画

中国的将来城市, 科幻插画

现代修建, 白雪纷飞

螺蛳粉

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

想体验 OneFlow 版的太乙 Stable Diffusion？只须要批改两行代码：

之所以能这么轻松迁徙模型，是因为 OneFlow Stable Diffusion 有两个杰出的个性：

OneFlowStableDiffusionPipeline.from_pretrained 可能间接应用 PyTorch 权重。
OneFlow 自身的 API 和 PyTorch 对齐，因而 import oneflow as torch 之后，torch.autocast、torch.float16 等表达式齐全不须要批改。

上述个性使得 OneFlow 兼容了 PyTorch 的生态，这不仅在 OneFlow 对太乙 Stable Diffusion 的迁徙中施展了作用，也大大减速了 OneFlow 用户迁徙其它许多模型，比方在和 torchvision 对标的 flowvision 中，许多模型只需通过在 torchvision 模型文件中退出 import oneflow as torch 即可失去。

此外，OneFlow 还提供全局“mock torch”性能，在命令行运行 eval $(oneflow-mock-torch) 就能够让接下来运行的所有 Python 脚本里的 import torch 都主动指向 oneflow。

深度学习算法原型开发阶段须要疾速批改和调试，动态图执行（Eager mode, define by run）最优。但在部署阶段，模型曾经固定下来，计算效率变得更重要，动态图执行（Lazy mode，define and run）能够借助编译器做动态优化来取得更好的性能。因而， 推理阶段次要应用动态图模式。

最近，PyTorch 降级到 2.0 引入了 compile() 这个 API，能够把一个模型或一个 Module 从动态图执行变成动态图执行。OneFlow 里也有一个相似的机制，不过接口名是 nn.Graph()，它能够把传入 Module 转成动态图执行模式。

不仅如此，OneFlow 的 nn.Graph 模式基于 MLIR 实现了一系列计算图的图层优化 ，譬如内存布局、算子交融等。

这不仅使得计算图示意的深度学习模型能够在各种硬件上达到最高性能，更重要的是，使得深度学习框架导入的计算图更不便地在不同硬件之间实现迁徙，有助于克服国产硬件软件生态单薄的问题。将来，咱们将公布更多内容来揭示 OneFlow 深度学习编译器的设计和实现。

欢送 Star、运行 OneFlow 版太乙 Stable Diffusion： https://github.com/Oneflow-In…

OneFlow 地址： https://github.com/Oneflow-In…

欢送 Star、试用 OneFlow 最新版本：https://github.com/Oneflow-In…

正文完