🎆

Raspberry Pi Pico 2 の DSP命令をMNIST+MLPで試す！

阿久津

2024/12/13に公開

公式のExamplesにDSPのサンプルが無かったので試してみました！

結論としては、DSPを使うとq1.15の積和演算能力が約5倍になっていることが確認できました！

比較結果表はこちらです。
最初はPico(無印)との全体的なスペック比較について記載してます。本題はDSP命令を試すからです。

 Raspberry Pi Pico 2 vs Raspberry Pi Picoそもそも2は何が違うのか？というところからです。

公式資料だと以下のように比較されています。基本スペック向上に加えて、ARM TrustZoneやRISV-V対応など色々と追加されています。
Dattasheetより
なんと言ってもCortex-M0+ ⇒ Cortex-M33です。演算関係だとFPUとDSPが使用可能になっている点は外せません。
公式のアーカイブ?より
さらに、こちらの記事では消費電力の減少が報告されており、Pico 2圧勝な感じです。
これが1000円程度なのは安いなぁと思います！

 DSPとはDSPはデジタルシグナルプロセッサの略称ですが、正式名称を見たところで何も分かりませんよね。
私のざっくり解釈です。DSPは高速積和演算器とよく言われますが、そもそも積はコストが大きい演算です。積というのは和を繰り返し行って計算します。したがって、これを汎用CPUに搭載されているレジスタベースの演算で直列実行すると時間が掛かります。これを並列で一気に実行できるのがDSPという理解です。
積和演算
 y = \sum_{i=1}^{n} x_i \times w_i 詳細はWikiへ！
https://ja.wikipedia.org/wiki/デジタルシグナルプロセッサ

 DSP命令を試す標題に戻ります。

 方針DSP命令と言ってもアゼンブリではなく、ARMの公開しているC Interfaceを呼ぶこととします。
数字の手書き画像データ(MNIST)でMLPを学習させ、分類タスクを解かせます。

そして分類処理が回っているところまで確認し、DSPの有無による実行時間への影響を確認することを目的とします。

あくまでDSP命令を試すのが本題なので、量子化などは手抜きでやります。
また、学習済みモデルをDSP命令へ自動変換可能なツールがいくつかありますが、折角なので単純明快に直接DSP命令を呼びます。

 登場人物
 モデルとデータモデルは最小限のMLPです。

後々の型変換のため、推論時にのみ動作するassertが入っています。
https://github.com/akutsu-kei/hello_dsp/blob/main/src_py/model.py
データとしては、torchvisionのMNISTを使用します。
特筆点は、重みと入出力データを全てDSP演算用の型に変換する必要がある点です。

その型とはq1.15(=q15)固定小数です。一般的な整数型をそのまま負の累乗として小数に拡張したような型です。浮動小数よりもシンプルです。


フォーマット
ビット構成
値域


q1.15
1ビット符号 + 15ビット小数
-1.0 ~ 0.999969482421875

変換方法は後述します。

 命令セットArm Cortex系の命令セットは、Common Microcontroller Software Interface Standard(CMSIS)としてまとめられています。
https://github.com/ARM-software/CMSIS_5
今回扱うDSP命令は、CMSIS-DSPとして以下で管理されています。CMSIS-DSPは↑のCMSIS_5/CORE Interfaceに依存しています。
https://github.com/ARM-software/CMSIS-DSP
最新のドキュメントはこちらです。Googleだと旧版の方が上位に出てくるので注意です。
https://arm-software.github.io/CMSIS-DSP/latest/
今回取り扱うMLPではBiasを入れないので、計算要素は行列積とRELUだけです。

行列積の計算にCMSIS-DSPの以下の命令を使用します。RELUは自前でやります。
arm_mat_mult_fast_q15()
arm_mat_mult_q15()
https://arm-software.github.io/CMSIS-DSP/latest/group__MatrixMult.html
(q1.15を選んだのはなんとなくです。)

 手順ファイル構造は以下のようになっています。
src_py: pythonを用いた学習・Cヘッダ出力関係をまとめています。
hello_dsp_q16.c: Cのmain関数コードです。arm_mat_mult_q15を呼ぶ推論関数もここです。
image_data.h, model_parameters.h: 実行用データ、MLP重みです。
extern: CMSIS_5, CMSIS_DSPをsubmoduleとしてリンクしています。
https://github.com/akutsu-kei/hello_dsp

 学習何てことはない通常の学習です。

mlp_mnist.pthファイルとして学習結果を出力します。
https://github.com/akutsu-kei/hello_dsp/blob/main/src_py/train.py

 推論テストと重みのファイル出力学習したモデルで推論テストをしてみます。

mlp_mnist.pthファイルを読み込んで推論してみると、各レイヤの変数値がq1.15の値域[-1, 1) をオーバーするため、以下のassertでエラーが発生します。

（assertを外せば推論自体は上手くできます。）
# akutsu-kei/hello_dsp/src_py/test_export.py
    def forward(self, x):
        if not self.training:
            assert torch.all(x > -1.1) and torch.all(x < 1.1)
        x = x.view(-1, 28 * 28)
        if not self.training:
            assert torch.all(x > -1.1) and torch.all(x < 1.1)
        x = self.fc1(x)
        if not self.training:
            assert torch.all(x > -1.1) and torch.all(x < 1.1)
        x = self.relu1(x)
        if not self.training:
            assert torch.all(x > -1.1) and torch.all(x < 1.1)
        x = self.fc2(x)
        return x
そこで、入力する画像を 1/30 にスケールして値を小さくします。各レイヤの変数値がq1.15の値域[-1, 1) をオーバーしなくなり、先ほどのassertを回避できます。（この状態でも推論自体は上手くできました。）
#src_py/transform.py
transform_test = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
    transforms.Lambda(lambda x: x / 30)
])
そして、重みについては運よくq1.15の値域[-1, 1) に収まっていることが確認できました。

値域[-1, 1)のassert通過後に、その重みをCヘッダーとして書き出しています。
def export_test_images(output_path):
    model = MLP()
    model.load_state_dict(torch.load('mlp_mnist.pth'))
    model.eval()

    scale_factor = 2**15

    with open(output_path, 'w') as f:
        f.write("#include <arm_math.h>\n\n")
        f.write(f"#define INPUT_SIZE {model.fc1.in_features}\n")
        f.write(f"#define FC1_SIZE {model.fc1.out_features}\n\n")
        f.write(f"#define FC2_SIZE {model.fc2.out_features}\n\n")

        for name, param in model.named_parameters():
            data = param.data.transpose(0, 1)
            assert torch.min(data) > -1.0　# ここを通過
            assert torch.max(data) < 1.0   # ここを通過
            data = data * scale_factor
            data = torch.clamp(data, -scale_factor, scale_factor - 1)
            name_fixed = name.replace('.', '_')
            f.write(f"const q15_t {name_fixed}[] = {{\n\t")
            data_list = [f"{round(x.item())}" for x in data.flatten()]
            f.write(", ".join(data_list))
            f.write("\n};\n\n")
この辺り、実際には１つのレイヤ毎にゼロ点とスケールファクタ、レンジを調整しながら、q1.15の値域を最大限に利用するのが通常です。こんなやり方でどうにかなるのは奇跡的かと思います。

 実行用の画像のファイル出力先ほどテストでも使用した10枚の画像をq1.15の値域[-1, 1]に収まる 1/30 にスケールしてCヘッダーとして書き出します。
https://github.com/akutsu-kei/hello_dsp/blob/main/src_py/export_images.py
ここで注意なのですが、q1.15は内部的には[-1,1)として扱われますが、表記としてはint16で扱います。
https://github.com/akutsu-kei/hello_dsp/blob/main/image_data.h

 Cのmain関数Cのmainコードは以下の通りです。

無限ループ内で10枚の画像を推論・テストしています。
// hello_dsp_q16.c
int main() {
    stdio_init_all();

    const uint LED_PIN = 16;
    gpio_init(LED_PIN);
    gpio_set_dir(LED_PIN, GPIO_OUT);

    while(true) {
        gpio_put(LED_PIN, 1);
        int index_of_maximum;
        pipeline(image_0, &index_of_maximum);
        assert(index_of_maximum == 0);
        pipeline(image_1, &index_of_maximum);
        assert(index_of_maximum == 1);
        pipeline(image_2, &index_of_maximum);
        assert(index_of_maximum == 2);
        pipeline(image_3, &index_of_maximum);
        assert(index_of_maximum == 3);
        pipeline(image_4, &index_of_maximum);
        assert(index_of_maximum == 4);
        pipeline(image_5, &index_of_maximum);
        assert(index_of_maximum == 5);
        pipeline(image_6, &index_of_maximum);
        assert(index_of_maximum == 6);
        pipeline(image_7, &index_of_maximum);
        assert(index_of_maximum == 7);
        pipeline(image_8, &index_of_maximum);
        assert(index_of_maximum == 8);
        pipeline(image_9, &index_of_maximum);
        assert(index_of_maximum == 9);
        gpio_put(LED_PIN, 0);
        sleep_ms(2000);
    }

    return 0;
}

 Cのforward関数特殊なことはしていませんが、一応は今回のキーポイントと思いますので。

arm_mat_mult_fast_q15とarm_mat_mult_q15とが切り替えられるようにしています。
// hello_dsp_q16.c
arm_status mlp_forward(const q15_t *input, q15_t *fc2_output_q15) {
    arm_status status;

    arm_matrix_instance_q15 mat_input;
    mat_input.numRows = 1;
    mat_input.numCols = INPUT_SIZE;
    mat_input.pData = (q15_t *)input;

    arm_matrix_instance_q15 mat_fc1_weights_instance;
    mat_fc1_weights_instance.numRows = INPUT_SIZE;
    mat_fc1_weights_instance.numCols = FC1_SIZE;
    mat_fc1_weights_instance.pData = (q15_t *)fc1_weight;

    arm_matrix_instance_q15 mat_fc1_output_instance;
    mat_fc1_output_instance.numRows = 1;
    mat_fc1_output_instance.numCols = FC1_SIZE;
    mat_fc1_output_instance.pData = fc1_output_q15;

    q15_t mat_mult_state_fc1[FC1_STATE_SIZE];
    #ifdef FAST_CMSIS_INSTRUCTION
    status = arm_mat_mult_fast_q15(&mat_input, &mat_fc1_weights_instance, &mat_fc1_output_instance, mat_mult_state_fc1);
    #else
    status = arm_mat_mult_q15(&mat_input, &mat_fc1_weights_instance, &mat_fc1_output_instance, mat_mult_state_fc1);
    #endif
    if(status != ARM_MATH_SUCCESS) {
        return status;
    }

    relu_q15(fc1_output_q15, FC1_SIZE);

    arm_matrix_instance_q15 mat_fc2_weights_instance;
    mat_fc2_weights_instance.numRows = FC1_SIZE;
    mat_fc2_weights_instance.numCols = FC2_SIZE;
    mat_fc2_weights_instance.pData = (q15_t *)fc2_weight;

    arm_matrix_instance_q15 mat_fc2_output_instance;
    mat_fc2_output_instance.numRows = 1;
    mat_fc2_output_instance.numCols = FC2_SIZE;
    mat_fc2_output_instance.pData = fc2_output_q15;

    q15_t mat_mult_state_fc2[FC2_STATE_SIZE];
    #ifdef FAST_CMSIS_INSTRUCTION
    status = arm_mat_mult_fast_q15(&mat_fc1_output_instance, &mat_fc2_weights_instance, &mat_fc2_output_instance, mat_mult_state_fc2);
    #else
    status = arm_mat_mult_q15(&mat_fc1_output_instance, &mat_fc2_weights_instance, &mat_fc2_output_instance, mat_mult_state_fc2);
    #endif
    if(status != ARM_MATH_SUCCESS) {
        return status;
    }

    return ARM_MATH_SUCCESS;
}

 推論時間 計測方法いよいよ計測していきましょう。

推論の始点・終点でGPIOピンの切替えて、その信号幅をオシロで観測し推論時間を計測します。
まずは、計測条件を決めます。
基本的には、メインのCMakeLists.txtの定義で計測条件を切り替えています。
https://github.com/akutsu-kei/hello_dsp/blob/main/CMakeLists.txt
DSP命令のスイッチ

FAST_CMSIS_INSTRUCTIONのマクロを切り替えることでDSPの無効/有効を切り替えます。
// hello_dsp_q16.c
#ifdef FAST_CMSIS_INSTRUCTION
status = arm_mat_mult_fast_q15(&mat_input, &mat_fc1_weights_instance, &mat_fc1_output_instance, mat_mult_state_fc1);
#else
status = arm_mat_mult_q15(&mat_input, &mat_fc1_weights_instance, &mat_fc1_output_instance, mat_mult_state_fc1);
#endif
最適化のスイッチ

PICO_DEOPTIMIZED_DEBUG=OFFで最適化ON

PICO_DEOPTIMIZED_DEBUG=ONで最適化OFF　です。
DSPのスイッチ
__ARM_FEATURE_DSPのマクロを0/1に切り替えることでDSPの無効/有効を切り替えます。
実装から辿りましたが、以下に記載がありました。
https://developer.arm.com/documentation/dui0774/g/chr1383660321827

 推論時間 計測結果結果は表のとおりです。


No.
DSP命令
最適化
DSP
推論時間[ms/10回]


1
arm_mat_mult_fast_q15
ON
有効
86ms

2
arm_mat_mult_q15
ON
有効
89ms

3
arm_mat_mult_fast_q15
ON
無効
94ms

4
arm_mat_mult_q15
ON
無効
455ms

5
arm_mat_mult_fast_q15
OFF
有効
86ms

6
arm_mat_mult_q15
OFF
有効
88ms

7
arm_mat_mult_fast_q15
OFF
無効
94ms

8
arm_mat_mult_q15
OFF
無効
455ms

結果より分かることとして
PICO_DEOPTIMIZED_DEBUGの設定は推論時間に寄与しない。
arm_mat_mult_fast_q15はDSP有効、無効に関係なく高速。arm_mat_mult_fast_q15は精度を犠牲に高速に演算可能なように実装されているとのこと。アルゴリズムで演算数を減らしているようですが、速すぎます。調査の余地ありです。
DSP有無の差が出たのはarm_mat_mult_q15です。有効すると約5倍ほど高速に演算可能でした。

ぐらいかなと思います。
またPICO(無印)との比較を考えると、正確なq1.15行列積演算が約5倍速くできるようになった！、という感じです。
参考までにMAC演算数は１回のForwardでおおよそ、50,816です。
  784 \times 64 + 64 \times 10 = 50,816 No.6から見るに、DSPの演算能力はこの値以上はあるのかなと思います。
 50,816 \times 10[回]  / 0.088[sec]　/  1,000,000 = 5.77 MMACs/s No.1の計測画像のみ張っておきます。


 終わりにPICO2すごく良いです。

高速になってはいますが、価格は安く、省電力です。
今後としては、カメラと接続して画像認識させたり、Pytorchの重みをそのままCに移植してarm_mat_mult_f32命令を試したり、arm_mat_mult_fast_q15のアルゴリズムを覗いたり、色々とネタはありそうです。

 Tipデバッガがたまーに固まります。F1で以下のコマンドを実行すると、解消できました。

フォーマット	ビット構成	値域
q1.15	1ビット符号 + 15ビット小数	-1.0 ~ 0.999969482421875

No.	DSP命令	最適化	DSP	推論時間[ms/10回]
1	arm_mat_mult_fast_q15	ON	有効	86ms
2	arm_mat_mult_q15	ON	有効	89ms
3	arm_mat_mult_fast_q15	ON	無効	94ms
4	arm_mat_mult_q15	ON	無効	455ms
5	arm_mat_mult_fast_q15	OFF	有効	86ms
6	arm_mat_mult_q15	OFF	有効	88ms
7	arm_mat_mult_fast_q15	OFF	無効	94ms
8	arm_mat_mult_q15	OFF	無効	455ms

ヘッドウォータース

株式会社ヘッドウォータースのテックブログです。生成AI、LLM、Azureのサービスや資格、IoT、XR系などData&AIとApp modernizeに関して幅広く投稿します！

Raspberry Pi Pico 2 の DSP命令をMNIST+MLPで試す！

Raspberry Pi Pico 2 vs Raspberry Pi Pico

DSPとは

DSP命令を試す

方針

登場人物

モデルとデータ

命令セット

手順

学習

推論テストと重みのファイル出力

実行用の画像のファイル出力

Cのmain関数

Cのforward関数

推論時間計測方法

推論時間計測結果

終わりに

Tip

Discussion

Raspberry Pi Pico 2 vs Raspberry Pi Pico

DSPとは

DSP命令を試す

方針

登場人物

モデルとデータ

命令セット

手順

学習

推論テストと重みのファイル出力

実行用の画像のファイル出力

Cのmain関数

Cのforward関数

推論時間 計測方法

推論時間 計測結果

終わりに

Tip

Discussion

推論時間計測方法

推論時間計測結果