目的

Module-LLMが搭載しているaxera-techのax630cは、CPUに2-coreのCoretex-A53を搭載していることに加えて、3.2Tops @int8のNPUを搭載しています。Module-LLMで、PythonにてCPU推論とNPU推論を行ったときの速度を比べてみました。

結果

Module-LLMのresnetモデルで、onnxruntimeを使ったCPU推論とNPU推論の速さの違いのグラフを描いてみると以下のような結果となりました。CPU推論では、onnxruntimeの機能を使って、元のFloat32の推論のものと、INT8量子化を行った推論のものを掲載しています。

Module-LLMで色々なクラス分類のモデルでonnxruntime(CPU)とpyaxengine(NPU)の速度を色々比べてみた。やはり、NPUは早い…30倍～100倍くらい早い感じだろうか。 pic.twitter.com/Mk7i6nZ9X3
— nnn (@nnn112358) December 15, 2024

注意事項

　今回は、Pythonでの評価を行いましたが、PythonよりもC++を使った方が、推論は早くなることがわかっています。また、axengineもまだ実装が行われているため、今後、評価値が変化する可能性があります。
よって、この結果は、Module-LLMのベストエフォートを示すものではありません。

ONNXモデルのエスクポートとAXモデルへの変換

UbuntuPCを使い、Pytorchのクラス分類のモデルから、ONNXモデルへのエスクポートを行いました。次に、onnxruntimeの機能を使って、INT8量子化を行いました。
更に、ONNXモデルにAxera-techのモデル変換ツールpulsar2を使って、Module-LLMのNPUが対応するAXモデルへの変換を行いました。

$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
$ cd model_conv
$ mkdir -p models
$ mkdir -p quantized_models
$ mkdir -p calibration_images
$ cd models
$ python classification_torch_model_export.py
$ cd ..
$ python onnx_quantize_static_image_batch.py

Module-LLMでの実行

変換したモデルをModule-LLMにコピーし、Module-LLMのNPUをPythonで扱う開発環境である、axengineをインストールしました。
ONNXモデルからonnxrutimeでCPU推論を行い、AXモデルからaxengineでNPU推論を行いました。

root@m5stack-LLM:# wget https://github.com/AXERA-TECH/pyaxengine/releases/download/0.0.1rc3/axengine-0.0.1-py3-none-any.whl
root@m5stack-LLM:# pip install axengine-0.0.1-py3-none-any.whl
root@m5stack-LLM:# ./batch_run_onnx.sh
root@m5stack-LLM:# ./batch_run_axmodel.sh

axengineは、onnxrutimeと同じく、推論をsession.run()関数から呼び出すことができます。そのため、onnxrutimeからの移植が大変容易です。

実行した結果とソースコードは、以下のリポジトリに掲載しています。