OnnxStream: 轻量级ONNX推理引擎助力边缘设备运行大型AI模型

OnnxStream

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>();
// 处理推理结果...