Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

把Stable Diffusion模型塞进iPhone里,做成APP一分钟出图

Stable Diffusion 可能很快就会在移动端普及。

在 iPhone 上运行 Stable Diffusion 到底难不难?今天我们要介绍的这篇文章,作者给出了答案:不难,而且 iPhone 还剩余 50% 的性能。

众所周知,每年苹果都会推出一款声称在各方面都更快、更好的新款 iPhone,这主要得益于新的视觉模型和图像传感器的快速发展。就拿拍照来说,如果回到 10 年前,你能用 iPhone 拍出高质量的图片吗,答案是不能,因为技术的发展是渐进式的,10 年时间,足够提高手机拍照技术。

由于技术的这种发展模式(渐进式),在一段时间里有些程序即使运行在最好的计算设备上,也几乎无法使用。但是这些带有新启用场景的新程序吸引了一些用户的注意力,人们愿意研究它。

本文的作者就是被吸引的其中之一,在过去 3 周里,作者开发了一个应用程序,可以通过 Stable Diffusion 来生成(summon)图像,然后按你喜欢的方式编辑它。该应用在最新的 iPhone 14 Pro 上生成图片仅需一分钟,使用大约 2GiB 的应用内存,另外还需要下载大约 2GiB 的初始数据才能开始使用。

应用商店链接:https://apps.apple.com/us/app/draw-things-ai-generation/id6444050820

这一结果引来众多网友讨论,有人开始担心手机耗电问题,并开玩笑的说:这很酷,不过这看起来是一个消耗手机电池的好方法。

图片

「我从来没有像现在这样开心地感受 iPhone 的热度。」

图片

「这个寒冬,可以把手机当暖手器用了。」

图片

不过在大家调侃手机发热问题的同时,他们也给与这项工作极高的评价。

「这简直不可思议。在我的 iPhone SE3 上生成一张完整的图像大约需要 45 秒——这几乎和我的 M1 Pro macbook 用原始版本生成的速度一样快!」

图片

内存、硬件同时优化

这是如何做到的呢?接下来我们看看作者的实现过程:

想要完成在 iPhone 上运行 Stable Diffusion,还能结余 50% 的性能,面临的一大挑战是需要在 6GiB RAM 的 iPhone 设备上将程序运行起来。6GiB 听起来很多,但如果你在 6GiB 设备上使用超过 2.8GiB,或在 4GiB 设备上使用超过 2GiB,iOS 就会杀死你的应用程序。

那么 Stable Diffusion 模型究竟需要多少内存来进行推理?

这还要从模型的结构说起。通常 Stable Diffusion 模型包含 4 个部分:1. 文本编码器,它生成文本特征向量以指导图像生成;2. 可选的图像编码器,将图像编码到潜在空间 (用于图像到图像生成);3. 降噪器模型,它从噪声中缓慢地去噪图像的潜在表示;4. 图像解码器,从潜在表示中解码图像。

第 1、第 2 和第 4 个模块在推理过程中运行一次,最大需要约 1GiB。而降噪器模型占用了大约 3.2GiB(全浮点数),并且还需要执行多次,因此作者想让该模块在 RAM 中保存得更久。

最初的 Stable Diffusion 模型需要接近 10GiB 才能执行单个图像推理。在单个输入(2x4x64x64)与输出(2x4x64x64)之间,其中夹杂着许多输出层。并不是所有层的输出都可以被立即复用,它们中一部分必须保留一些参数以供后续使用(残差网络)。

一段时间以来,研究者围绕 PyTorch Stable Diffusion 进行了一番优化,对 PyTorch 用到的 NVIDIA CUDNN 和 CUBLAS 库,他们保留了暂存空间,这些优化都是为了降低内存使用量,因此 Stable Diffusion 模型可以用低至 4GiB 的卡运行。

但这仍然超出了作者的预期。因此作者开始专注于苹果硬及优化。

起初作者考虑的是 3.2GiB 或 1.6GiB 半浮点数,如果不想触发苹果的 OOM(Out of Memory,指的是 App 占用的内存达到了 iOS 系统对单个 App 占用内存上限后,而被系统强杀掉的现象),作者大约有 500MiB 的空间可以使用。

第一个问题,每个中间输出的大小到底是多少?

事实证明,它们中的大多数都相对较小,每个都低于 6MiB (2x320x64x64)。作者使用的框架 (s4nnc) 可以合理地将它们打包到小于 50MiB,以备复用。

值得一提的是,降噪器有一个自注意机制,它以自己的图像潜在表示作为输入。在自注意力计算期间,有一个大小为 16x4096x4096 的批处理矩阵,对该矩阵应用 softmax 后,大约是 FP16 中的 500MiB,并且可以「inplace」完成,这意味着它可以安全地重写其输入而不会损坏。幸运的是,Apple 和 NVIDIA 低级库都提供了 inplace softmax 实现,然而 PyTorch 等更高级的库中没有。

那么是否真的使用 550MiB + 1.6GiB 左右的内存就能完成?

在 Apple 硬件上,实现神经网络后端的一个常用选择是使用 MPSGraph 框架。于是作者首先尝试使用 MPSGraph 实现了所有的神经网络操作。在 FP16 精度下峰值内存使用量大约是 6GiB,显然比预期的内存使用量多太多,这是怎么回事?

作者详细分析了原因,首先他没有按照常见的 TensorFlow 方式使用 MPSGraph。MPSGraph 需要对整个计算图进行编码,然后使用输入 / 输出张量,进而处理内部分配,并让用户提交整个图以供执行。

而作者使用 MPSGraph 的方式很像 PyTorch 的做法——当作一个操作执行引擎。为了执行推理任务,许多已编译的 MPSGraphExecutable 在 Metal 命令队列上执行,它们中的每一个都可能持有一些中间分配内存。如果一次性提交,那么所有这些命令都持有分配内存,直到完成执行。

一种解决这个问题的简单方法是调整提交速度,没有必要一次性提交所有命令。实际上,Metal 的每个队列有 64 个并发提交的限制。作者尝试改成一次提交 8 个操作,峰值内存就降低到了 4GiB。

然而,这仍然比 iPhone 能承受的多 2 GiB。

为了使用 CUDA 计算自注意力,原始 Stable Diffusion 代码实现中有一个常见技巧:使用置换而不是转置。这个技巧很有效,因为 CUBLAS 可以直接处理置换的跨步(strided)张量,避免使用专用内存来转置张量

但是 MPSGraph 没有跨步张量支持,一个置换的张量无论如何都会在内部被转置,这需要中间分配内存。通过显式转置,分配将由更高级别的层处理,避免了 MPSGraph 内部效率低下。利用这个技巧,内存使用量将接近 3GiB。

事实证明,从 iOS 16.0 开始,MPSGraph 不能再为 softmax 做出最优分配决策。即使输入和输出张量都指向相同的数据,MPSGraph 也会分配一个额外的输出张量,然后将结果复制到指向的位置。

作者发现使用 Metal Performance Shaders 替代方案完全符合要求,并将内存使用量降至 2.5GiB,而不会出现任何性能下降。

另一方面,MPSGraph 的 GEMM 内核需要内部转置。显式转置在此也无济于事,因为这些转置不是更高级别层的「inplace」操作,对于特定的 500MiB 大小的张量,这种额外的分配是不可避免的。通过切换到 Metal Performance Shaders,项目作者又回收了 500MiB,性能损失约为 1%,最终将内存使用量减到了理想的 2GiB。

参考链接:

https://news.ycombinator.com/item?id=33539192

https://liuliu.me/eyes/stretch-iphone-to-its-limit-a-2gib-model-that-can-draw-everything-in-your-pocket/

工程Stable Diffusion
相关数据
自注意力技术

自注意力(Self-attention),有时也称为内部注意力,它是一种涉及单序列不同位置的注意力机制,并能计算序列的表征。自注意力在多种任务中都有非常成功的应用,例如阅读理解、摘要概括、文字蕴含和语句表征等。自注意力这种在序列内部执行 Attention 的方法可以视为搜索序列内部的隐藏关系,这种内部关系对于翻译以及序列任务的性能非常重要。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

TensorFlow技术

TensorFlow是一个开源软件库,用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品,如语音识别、Gmail、Google 相册和搜索,其中许多产品曾使用过其前任软件DistBelief。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

神经网络技术

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

图像生成技术

图像生成(合成)是从现有数据集生成新图像的任务。

5G技术

第五代移动通信系统(5th generation mobile networks),简称5G,是4G系统后的延伸。美国时间2018年6月13日,圣地牙哥3GPP会议订下第一个国际5G标准。由于物理波段的限制,5G 的网络也将会与其他通信技术并用,包含长距离的其他传统电信波段。

暂无评论
暂无评论~