稳定的扩散运行在Raspberry Pi上,260MB的RAM”持有”一个具有10亿个参数的大型模型。

干货分享8个月前更新 Youzhizhan
1,381 0


Secure Diffusion诞生于11个月前,它可以在消费者Gpu上运行的消息鼓舞了许多研究人员。不仅如此,苹果官方很快宣布将Stable Diffusion”塞”到iPhone、iPad和Mac中运行。这大大降低了Stable Diffusion对硬件设备的要求,使其逐渐成为人人都能用的”黑科技”。

现在,它甚至可以在Raspberry Pi Zero2上运行。

稳定的扩散运行在Raspberry Pi上,260MB的RAM图片

树莓Pi零2″一样小。 快五倍。」

这是什么概念?运行稳定的扩散不是一件容易的事。 它包含一个具有10亿个参数的大型变压器模型。 建议的最小RAM/VRAM通常为8gb。RPI Zero2只是一个具有512mb内存的微型计算机。

这意味着在RPI Zero2上运行稳定扩散是一个巨大的挑战。而且,在操作过程中,作者没有增加存储空间,也没有将中间结果卸载到磁盘。

一般来说,主要的机器学习框架和库专注于最小化推理延迟和/或最大化吞吐量,但以上所有这些都是以牺牲内存使用为代价的。因此,作者决定编写一个超小,可解锁的推理库,致力于最大限度地减少内存消耗。

OnnxStream做到了。

稳定的扩散运行在Raspberry Pi上,260MB的RAM

项目地址https://github.com/vitoplantamura/OnnxStream

OnnxStream基于将推理引擎与负责提供模型权重的组件解耦的想法,模型权重是从WeightsProvider派生的类。WeightsProvider的特化可以实现任何类型的模型参数加载、缓存和预取。例如,自定义WeightsProvider可以决定直接从HTTP服务器下载数据,而无需加载或将任何内容写入磁盘(这也是OnnxStream名称中存在流的原因)。有两个默认的WeightsProviders可用:DiskNoCache和DiskPrefetch。

与微软的推理框架OnnxStream相比,OnnxStream只需要消耗1/55的内存即可达到相同的效果,但速度(在CPU上)仅比前者慢0.5-2倍。

接下来您将看到稳定扩散在RPI Zero2上运行的效果,以及后面的方法it.It 需要注意的是,虽然运行速度较慢,但对于大型机型来说,在更小,更有限的设备上运行是一种全新的尝试。

稳定的扩散运行在Raspberry Pi上,260MB的RAM图片

网友觉得这个项目很酷

在Raspberry Pi Zero2上运行稳定扩散

VAE解码器是唯一一个稳定扩散的模型,不能放置在RPI零2RAM与单精度或半精度。这是因为模型中存在残差连接,非常大的张量和卷积。唯一的解决方案是静态量化(8位)。

以下图像由作者的repo中包含的稳定扩散示例实现在不同精度VAE解码器下使用OnnxStream生成。

第一个图像是在作者的PC上生成的,使用rpi Zero2生成的相同潜像。

稳定的扩散运行在Raspberry Pi上,260MB的RAM精度为W16A16的VAE解码器的生成效果

稳定的扩散运行在Raspberry Pi上,260MB的RAM精度为W8A32的VAE解码器的生成效果

第三张图片由RPI Zero2在大约3小时内生成。

稳定的扩散运行在Raspberry Pi上,260MB的RAM图注:精度为W8A8的VAE解码器的生成效果

OnnxStream的特点

  • 将推理引擎与WeightsProvider解耦
  • WeightsProvider可以是DiskNoCache、DiskPrefetch或custom
  • 注意切片
  • 动态量化(8位无符号、非对称、百分位数)
  • 静态量化(W8A8无符号,不对称,百分位数)
  • 轻松校准定量模型
  • 支持FP16(有或没有FP16操作)
  • 实现了24个ONNX运算符(最常用的运算符)
  • 操作按顺序执行,但所有运算符都是多线程的
  • 单实现文件+头文件
  • XNNPACK调用封装在XnnPack类中(用于将来的替换)

并且需要注意的是,OnnxStream依赖于XNNPACK来加速某些基元:MatMul,Convolution,element-Smart Add/Sub/Mul/Div,Sigmoid和Softmax。

性能比较

稳定扩散由三个模型组成:文本编码器(672个操作和1.23亿个参数),UNET模型(2050个操作和8.54亿个参数)和VAE解码器(276个操作和4900万个参数)。

假设批处理大小等于1,则需要10个步骤来生成完整的图像。 这需要运行文本编码器2次,运行UNET模型20次(即2*10),并运行VAE解码器1次以获得良好的结果(使用欧拉祖先调度器)。

该表显示了稳定扩散的三种模型的不同推理时间,以及内存消耗(即Windows中的峰值工作集大小或Linux中的最大常驻集大小)。

可以发现,在UNET模型中(以FP16精度运行,在OnnxStream中启用了FP16算术),OnnxStream的内存消耗仅为OnnxRuntime的1/55,但速度仅慢0.5-2倍。

本次测试需要注意的点是:

  • OnnxRuntime的第一次运行是预热推理,因为它的InferenceSession是在第一次运行之前创建的,并在所有后续运行中重用。OnnxStream没有预热推理,因为它的设计纯粹是”渴望”的(但是,后续运行可以从操作系统的权重文件缓存中受益)。
  • 目前OnnxStream不支持批量大小! =1输入,这与OnnxRuntime不同,OnnxRuntime在运行UNET模型时使用批处理大小=2,以大大加快整个扩散过程。
  • 在测试中,更改OnnxRuntime的SessionOptions(如EnableCpuMemArena和ExecutionMode)对结果没有显着影响。
  • 在内存消耗和推理时间方面,OnnxRuntime的性能与NCNN(另一个框架)非常相似。
  • 测试操作条件:Home windows Server2019,16gb内存,8750H CPU(AVX2),970EVO Plus SSD,VMware上的8个虚拟内核。

注意切片和量化

在运行UNET模型时,使用”注意力切片”技术,并对VAE解码器使用W8A8量化,这对于将模型的内存消耗降低到适合在RPI Zero2上运行的水平至关重要。

虽然互联网上有很多关于定量神经网络的信息,但关于”注意力切片”的信息却很少。

这里的想法很简单:目标是在计算UNET中各种多头注意力的缩放点积注意力时避免生成完整的Q@Ok^T矩阵model.In UNET模型中,当注意头的个数为8时,Q的形状为(8,4096,40),Ok^T为(8,40,4096)。因此,第一个MatMul的最终形状是(8,4096,4096),这是一个512MB张量(FP32精度)。

稳定的扩散运行在Raspberry Pi上,260MB的RAM图片

解决方法是垂直划分Q,然后对每个Q块进行正常的注意操作。Q_sliced的形状为(1,x,40),其中x为4096(在此示例中),除以onnxstream::Mannequin::m_attention_fused_ops_parts(默认值为2,但可以自定义。

这个简单的技巧可以降低UNET模型在FP32精度从1.1gb到300mb运行时的整体内存消耗。更有效的替代方案是使用FlashAttention,但FlashAttention需要为每个支持的体系结构(AVX,NEON)等编写自定义内核。,在作者给出的例子中绕过了XnnPack。

有关详细信息,请参阅项目的GitHub界面。

[ad]
© 版权声明

相关文章

暂无评论

您必须登录才能参与评论!
立即登录
暂无评论...