vitoplantamura/OnnxStream

何を解決するか

OnnxStream は、大規模なAIモデル（例：Stable Diffusion や LLM）を極めて限られたRAMを持つハードウェア（例：512MBのRAMを持つRaspberry Pi Zero 2）で実行できるように設計された軽量な推論ライブラリです。標準的な機械学習フレームワークは速度とスループットを最優先するため、低価格端末が備えるメモリ（数ギガバイト）を必要とする場合が多く、そのようなデバイスでは実行が困難です。

動作方法

このライブラリは WeightsProvider を介して推論エンジンと重みの読み込みプロセスを分離することで、メモリ消費を最小限に抑えます。これにより、重みをRAMに完全にロードするのではなく、ディスクやHTTPサーバーからストリーミングで読み込むことが可能になります。以下のメモリ節約技術を採用しています：

• Attention Slicing：マルチヘッドアテンション中に巨大な中間テンソルが生成されるのを防ぐために、クエリテンソルを小さなチャンクに分割します。
• Quantization：動的（8ビット符号なし）および静的（W8A8）の両方の量子化をサポートし、重みと活性化の精度とサイズを低減します。
• Tiled Decoding：大規模なVAEデコーダ（SDXLなど）に対して、入力テンソルを重複するタイルに分割し、個別にデコードしてから再結合します。
• XNNPACK統合：MatMulやConvolutionなどの加速プリミティブにXNNPACKを使用します。

対象ユーザー

• 超低リソースの組み込みデバイス（例：Raspberry Pi Zero）をターゲットとする開発者。
• 専用GPUや高メモリを備えないハードウェア上で大規模な生成モデルを実行したいユーザー。
• ハッキング可能で最小限のC++推論エンジンを必要とするエンジニア（PythonおよびC#バインディング付き）。

特徴

• 極めて高いメモリ効率：特定のモデルではOnnxRuntimeと比較して最大55倍のメモリ削減が可能。
• 広範なモデル対応：Stable Diffusion 1.5、SDXL 1.0、SDXL Turbo、TinyLlama、Mistral 7B、YOLOv8、OpenAIのWhisperを実行可能。
• クロスプラットフォーム：Linux、Mac、Windows、Termux、FreeBSD、WebAssembly（ブラウザベース実行）をサポート。
• 柔軟な重み読み込み：RAM、ディスク（プリフェッチあり/なし）、カスタムストリーミングプロバイダーから重みを提供可能。