在iPhone上运行StableDiffusion很难吗？在我们今天要介绍的这篇文章中，作者给出了答案：不难，iPhone还有50%的性能剩余。众所周知，苹果每年都会推出一款声称在各个方面都更快更好的新iPhone，这主要得益于新视觉模型和图像传感器的快速发展。拍照，如果回到10年前，iPhone能不能拍出高质量的照片，答案是否定的，因为科技的发展是渐进的，10年足以提升手机拍照的技术.由于这种技术发展模式（增量），有时某些程序即使在最好的计算设备上也几乎无法使用。但是这些新启用场景的新程序已经引起了一些用户的关注，人们也愿意去研究。这篇文章的作者就是其中之一，在过去的3周里，作者开发了一款应用程序，可以通过StableDiffusion生成（召唤）图像，然后根据需要进行编辑。该应用仅需一分钟即可在最新的iPhone14Pro上生成图像，使用约2GiB的应用内存，并需要下载约2GiB的初始数据才能开始使用。应用商店链接：https://apps.apple.com/us/app/draw-things-ai-generation/id6444050820这个结果引起了很多网友的讨论。有些人开始担心手机的耗电问题，并开玩笑说：这很酷，但它看起来像是一个耗尽手机电池的好方法。“我从未如此高兴地感受到iPhone的热量。”“这个寒冷的冬天，我可以用手机当暖手宝。”不过大家在拿手机发热问题开玩笑的同时，也给了这部作品极高的评价。“太不可思议了。在我的iPhoneSE3上生成完整图像大约需要45秒——这几乎和我的M1Promacbook用原始版本生成它一样快！”内存、硬件同时优化这是怎么回事你到哪儿了？接下来我们看看作者的实现过程：为了在iPhone上运行StableDiffusion并且仍然有50%的性能剩余，一个主要的挑战是在6GiBRAM的iPhone设备上运行该程序。6GiB听起来很多，但如果您在6GiB设备上使用超过2.8GiB，或在4GiB设备上使用超过2GiB，iOS将终止您的应用程序。那么StableDiffusion模型推理需要多少内存呢？这还要从模型的结构说起。通常StableDiffusion模型由4部分组成：1.Textencoder，生成文本特征向量，指导图像生成；2.可选的图像编码器，将图像编码到潜在空间（用于图像到图像的生成）；3..一个降噪器模型，它可以慢慢地从噪声中去除图像的潜在表示；4.图像解码器，从潜在表示中解码图像。模块1、2和4在推理期间运行一次，最多需要~1GiB。降噪器模型占用大约3.2GiB（全浮点数）并且需要执行多次，因此作者希望将此模块保留在RAM中更长时间。原始的StableDiffusion模型需要接近10GiB才能执行单图像推理。在单个输入（2x4x64x64）和输出（2x4x64x64）之间，散布着许多输出层。并非所有层输出都可以立即重用，其中一些必须保留一些参数供以后使用（残差网络）。一段时间以来，研究人员优化了PyTorchStableDiffusion。对于PyTorch使用的NVIDIACUDNN和CUBLAS库，它们预留了临时存储空间。这些优化都是为了减少内存使用，因此StableDiffusion模型可以在低至4GiB的卡上运行。但还是超出了作者的预期。于是笔者开始关注苹果硬件和优化。起初，作者考虑了3.2GiB或1.6GiB的半浮点数。如果不想触发苹果的OOM（OutofMemory，即App占用的内存已经达到iOS系统对单个App占用内存的上限，会被系统kill掉.现象），作者有大约500MiB的可用空间。第一个问题，每个中间输出的大小到底是多少？事实证明，它们中的大多数都比较小，每个都在6MiB(2x320x64x64)以下。作者使用的框架（s4nnc）可以合理地将它们打包成不到50MiB以供重用。值得一提的是，降噪器有一个自我注意机制，将它自己的图像潜在表示作为输入。自注意力计算时，有一个大小为16x4096x4096的batch矩阵，对这个矩阵应用softmax后，在FP16中约为500MiB，并且可以“就地”完成，这意味着它可以安全地重写其输入而不会被损坏.幸运的是，Apple和NVIDIA的低级库都提供了就地softmax实现，而PyTorch等高级库则没有。那么真的可以使用大约550MiB+1.6GiB的内存吗？在Apple硬件上实现神经网络后端的常见选择是使用MPSGraph框架。所以作者首先尝试使用MPSGraph来实现所有的神经网络操作。在FP16精度下，峰值内存使用量约为6GiB，显然内存使用量比预期多得多，这是怎么回事？笔者详细分析了原因。首先，他没有按照通常的TensorFlow方式使用MPSGraph。MPSGraph需要对整个计算图进行编码，然后使用输入/输出张量，处理内部分配，让用户提交整个图来执行。作者使用MPSGraph的方式与PyTorch的做法非常相似——作为运算执行引擎。为了执行推理任务，许多已编译的MPSGraphExecutables在Metal命令队列上执行，每个MPSGraphExecutables可能持有一些中间分配的内存。如果一次全部提交，所有这些命令都会保留分配的内存，直到它们完成执行。解决这个问题的一个简单方法就是调整提交速度，没必要一次提交所有命令。事实上，Metal对每个队列有64个并发提交的限制。笔者尝试一次提交8个操作，峰值内存降为4GiB。但是，这仍然比iPhone可以处理的多2GiB。要使用CUDA计算自注意力，在最初的稳定扩散代码实现中有一个常见的技巧：使用排列而不是转置。这个技巧之所以有效，是因为CUBLAS可以直接处理排列的跨步张量，避免使用专用内存来转置张量。但是MPSGraph没有跨步张量支持，排列张量无论如何都会在内部转置，这需要中间内存分配。通过显式转置，分配将由更高级别的层处理，避免MPSGraph内部的低效率。使用这个技巧，内存使用量将接近3GiB。事实证明，从iOS16.0开始，MPSGraph无法再为softmax做出最优分配决策。即使输入和输出张量都指向相同的数据，MPSGraph也会分配一个额外的输出张量，然后将结果复制到指向的位置。作者发现使用MetalPerformanceShaders替代方案完全符合要求，并且将内存使用量降至2.5GiB而没有任何性能影响。另一方面，MPSGraph的GEMM内核需要内部转置。显式转置在这里也无济于事，因为这些转置不是更高级别层的“就地”操作，并且对于500MiB的特定张量大小，这种额外分配是不可避免的。通过切换到MetalPerformanceShaders，项目作者又回收了500MiB，性能损失约1%，最终将内存使用量降低到理想的2GiB。
