OnnxStream:让边缘设备也能运行大型AI模型
在人工智能快速发展的今天,越来越多的设备需要具备运行复杂AI模型的能力。然而,受限于内存和计算资源,许多边缘设备难以部署大型模型。OnnxStream应运而生,它是一个专注于最小化内存消耗的轻量级ONNX推理库,让边缘设备也能运行Stable Diffusion等大型AI模型。
主要特性
OnnxStream具有以下主要特性:
- 推理引擎与权重提供组件解耦,支持灵活的权重加载方式
- 支持注意力切片(Attention Slicing)技术,大幅降低内存占用
- 支持动态量化和静态量化,进一步压缩模型大小
- 支持FP16精度推理
- 实现了39个最常用的ONNX算子
- 单文件实现,易于集成和定制
- 依赖XNNPACK进行底层加速
- 初步支持GPU加速(仅限LLM应用)
惊人的内存效率
OnnxStream在内存效率方面表现出色。以Stable Diffusion 1.5的UNET模型为例,OnnxStream的内存消耗比OnnxRuntime低55倍,同时延迟仅增加50%-200%。这使得在512MB RAM的树莓派Zero 2上运行Stable Diffusion XL成为可能。
支持多种模型
OnnxStream目前支持以下模型:
- Stable Diffusion 1.5
- Stable Diffusion XL 1.0 Base
- Stable Diffusion XL Turbo 1.0
- TinyLlama 1.1B
- Mistral 7B
在树莓派Zero 2上运行SDXL
通过一系列优化,OnnxStream成功在树莓派Zero 2上运行了Stable Diffusion XL,内存占用不到300MB。主要优化包括:
- 对UNET模型使用UINT8动态量化
- 对VAE解码器采用分块解码技术,将内存占用从4.4GB降至298MB
- 使用注意力切片技术,避免生成大型中间张量
在树莓派Zero 2上生成一张512x512的图像,SDXL Turbo仅需29分钟,而标准SDXL则需要11小时。
大语言模型支持
除了图像生成模型,OnnxStream还支持在桌面和服务器上运行TinyLlama 1.1B和Mistral 7B等大语言模型。通过使用RamWeightsProvider,OnnxStream可以将整个模型加载到内存中,实现交互式推理。
同时,OnnxStream还提供了初步的GPU加速支持。用户可以指定将部分模型卸载到GPU,实现CPU和GPU之间的并行计算。
灵活易用
OnnxStream提供了简单的C++ API,用户只需几行代码即可加载和运行模型:
Model model;
model.read_file("path_to_model_folder/model.txt");
tensor_vector<float> data;
// 填充输入数据...
Tensor t;
t.m_name = "input";
t.m_shape = { 1, 4, 64, 64 };
t.set_vector(std::move(data));
model.push_tensor(std::move(t));
model.run();
auto& result = model.m_data[0].get_vector<float>();
// 处理推理结果...
结语
OnnxStream为边缘设备运行大型AI模型开辟了新的可能性。通过精心的内存优化,它让树莓派Zero 2这样的微型设备也能运行Stable Diffusion XL等复杂模型。未来,OnnxStream还将继续优化性能,支持更多模型,为边缘AI的发展贡献力量。
