マイコンでディープラーニングした話 on ESP32 #Python

こんにちは。ESP32搭載のカメラキット、M5Camera を手に入れました。
ハイスペックマイコンESP32にカメラまで付いて来るとあっては、もうディープラーニング(推論)するしかないですね。
今回は画像認識の鉄板ネタ、MNISTデータセットを使った手書き数字認識に挑戦してみました。

ソースコード：https://github.com/yoku001/esp32-camera-dnn

手順

M5Cameraの公式サンプルをベースにディープラーニング周りの処理を実装していきます。

流石にESP32上でニューラルネットワークの学習まで行うのは困難なため、学習はPCなりクラウドなりで済ませたうえで、推論処理のみESP32の組込みOS(FreeRTOS) に移植する方針で進めて行きます。

悩ましいのはディープラーニングフレームワークの選定です。昨今のエッジコンピューティングブームもあってか、推論用にC/C++APIを提供するフレームワークは増えつつありますが、マイコンや組込みOS環境をターゲットにしたものは決して多くありません。

今回のデモ作成にあたっては、以下のフレームワークを検討しました。

Neural Network Libraries (nnabla)
- by SONY
- https://github.com/sony/nnabla
- GUIフロントエンドのNeural Network Consoleで有名な日本発DNNフレームワーク。公式のC言語ランタイムと組合せることで、学習済みモデルをCソースコードとしてデプロイ可能。
e-AIトランスレータ
- by Renesas Electronics
- https://www.renesas.com/jp/ja/solutions/key-technology/e-ai.html
- フレームワークではなくTensorflow, Caffeの学習済みモデルをCソースコードに変換するコンバータ。対応するネットワークはやや物足りないが、ランタイムが不要でスタンドアロンで動作するソースコードが出力される。
OpenCV dnn module
- by OpenCV team
- https://github.com/opencv/opencv
- 皆大好きOpenCV。多機能だがメモリフットプリント、バイナリサイズともかなりダイエットしないとM5Cameraに載せるのは難しそう。

今回は対応するネットワークの種類が多く、フットプリントもそれなりに小さいNeural Network Librariesを使用することにしました。

1. モデル学習

MNIST分類モデルの作成にはnnabla公式サンプルのCNNを使ったサンプルコードを使用しました。このサンプルにはLeNet風のネットワークとResNet風のネットワークが用意されていますが、今回はより単純なLeNet風ネットワークを使用していきます。

classification.py

def mnist_lenet_prediction(image, test=False, aug=None):
    """
    Construct LeNet for MNIST.
    """
    image /= 255.0
    image = augmentation(image, test, aug)
    c1 = PF.convolution(image, 16, (5, 5), name='conv1')
    c1 = F.relu(F.max_pooling(c1, (2, 2)), inplace=True)
    c2 = PF.convolution(c1, 16, (5, 5), name='conv2')
    c2 = F.relu(F.max_pooling(c2, (2, 2)), inplace=True)
    c3 = F.relu(PF.affine(c2, 50, name='fc3'), inplace=True)
    c4 = PF.affine(c3, 10, name='fc4')
    return c4

まずはパフォーマンスの事は考えず、そのまま学習を進めてみます。

git clone https://github.com/sony/nnabla-examples.git
cd nnabla-examples/mnist-collection
python classification.py -c cudnn -n lenet -o output

学習完了後、出力ディレクトリにnnabla形式のモデルファイルが出力されます。(上記の例であれば ./output/lenet_result.nnp )

2. Cソースへの変換

nnabla_cliでモデルファイルをCソースコードに変換します。

mkdir ./output_csrc
nnabla_cli convert -O CSRC -b 1 ./output/lenet_result.nnp ./output_csrc

変換に成功すると下記の6ファイルが出力されます。

.
└── output_csrc
    ├── GNUmakefile                  テスト用Makefile
    ├── Validation_example.c         テスト用コード
    ├── Validation_inference.c       モデル定義類
    ├── Validation_inference.h       ↑のヘッダ
    ├── Validation_parameters.c      重みパラメータ類
    └── Validation_parameters.h      ↑のヘッダ

推論に必要となるのは Validation_inference(.c|.h) と Validation_parameters(.c|.h) の計4ファイル。
これらをM5Cameraサンプルプロジェクトのcomponents下にコピーしましょう。

cp -r ./output_csrc/* ~/esp/esp32-camera-dnn/components/dnn/

前述の通り変換コードのコンパイルにはnnabla公式のC言語ランタイムが必要となります。今回はランタイムもcomponents下に配置しました。

cd ~/esp/esp32-camera-dnn/components/dnn/
git submodule add https://github.com/sony/nnabla-c-runtime.git

3. 実装

M5Camera公式サンプルのmain/main.cに処理を実装していきます。

カメラモジュール初期化パラメータ

MNIST向けにカラーフォーマットをグレースケールに、画素数も後ほど縮小することを見越しQQVGAに変更しておきます。

main.c

    .pixel_format = PIXFORMAT_GRAYSCALE, //YUV422,GRAYSCALE,RGB565,JPEG
    .frame_size = FRAMESIZE_QQVGA, //QQVGA-UXGA Do not use sizes above QVGA when not JPEG

    .jpeg_quality = 20, //0-63 lower number means higher quality
    .fb_count = 1 //if more than one, i2s runs in continuous mode. Use only with JPEG

推論

推論用のタスクを作成していきます。今回はM5CameraサンプルのJPEGストリーミング処理(jpg_stream_httpd_handler)に追記していく形で実装しました。

main.c

esp_err_t jpg_stream_httpd_handler(httpd_req_t *req) {
    // タスク初期化
    // ...

    float *nn_input_buffer = nnablart_validation_input_buffer(_context, 0);

    while (true) {
        // カメラ画像取得
        camera_fb_t *fb = esp_camera_fb_get();

        // ...

        // リサイズ
        stbir_resize_uint8(fb->buf, 160, 120, 0, resized_img, 28, 28, 0, 1);
        esp_camera_fb_return(fb);

        // 白黒反転, 閾値処理
        for (int i = 0; i < NNABLART_VALIDATION_INPUT0_SIZE; i++) {
            uint8_t p = ~(resized_img[i]);

            if (p < 180) {
                p = 0;
            }

            nn_input_buffer[i] = p;
        }

        // 推論実行
        nnablart_validation_inference(_context);

        // 推論結果の取得
        float *pred = nnablart_validation_output_buffer(_context, 0);
        int top_class = 0;
        float top_probability = 0.0f;
        for (int class = 0; class < NNABLART_VALIDATION_OUTPUT0_SIZE; class++) {
            if (top_probability < pred[class]) {
                top_probability = pred[class];
                top_class = class;
            }
        }

        // ...

        // 推論結果, 処理時間送信
        ESP_LOGI(TAG, "Result %d   Frame-time %ums (Inferrence-time %ums)",
            top_class, (uint32_t)frame_time, (uint32_t)infer_time);
    }

    // タスク終了
    nnablart_validation_free_context(_context);

    // ...
}

推論に関わる処理のみ取り上げると

推論用メモリの確保(nnablart_validation_input_buffer)
推論ループ
1. カメラ画像の取得
2. 前処理(画像のリサイズ、白黒反転、閾値処理。リサイズにはstb_imageライブラリを使用しています)
3. 推論実行(nnablart_validation_inference)
4. 推論結果の取得(nnablart_validation_output_buffer)
5. 推論結果の送信
推論用メモリの解放(nnablart_validation_free_context)

という感じになります。詳細は上記のgitリポジトリを参照してください。

推論結果はESP_LOG(USBシリアル出力)でPCに送信しています。ESP-IDFのMonitor機能で出力結果を確認できます。

make monitor PRINT_FILTER="camera"

M5CameraのAPに接続後、"192.168.4.1"へアクセスすると推論が開始されます。

4. デモ

左がUSBシリアルの出力、右がカメラ画像のストリーミングです。前処理が雑な事もあってか精度はイマイチですが、どうやら数字認識自体は上手く行っているようです。

気になる実行時間ですが、(前処理やシリアル通信時間も含めた)フレーム辺りの処理時間が約360ms、そのうち純粋に推論処理にかかる時間が約220msのようです。うーん...もうちょっと推論時間を短くしたいところですね。

推論の高速化と言ってもアプローチは種々ありますが、ここでは一番手っ取り早そうなニューラルネットワークモデル自体の軽量化に取り組んでみます。

高速化

元のLeNet風モデルはMNISTの分類にしてはかなり贅沢な構成になっていました。精度に影響しない範囲で軽量化したモデルがこちらです。

classification.py

def mnist_lenet_light_prediction(image, test=False, aug=None):
    """
    Construct LeNet for MNIST.
    """
    image /= 255.0
    image = augmentation(image, test, aug)
    c1 = PF.convolution(image, 4, (3, 3), name='conv1')
    c1 = F.relu(F.max_pooling(c1, (2, 2)), inplace=True)
    c2 = PF.convolution(c1, 16, (1, 1), name='conv2')
    c2 = PF.depthwise_convolution(c2, (3, 3), name='dw_conv2')
    c2 = F.relu(F.max_pooling(c2, (2, 2)), inplace=True)
    c3 = F.relu(PF.affine(c2, 30, name='fc3'), inplace=True)
    c4 = PF.affine(c3, 10, name='fc4')
    return c4

CNN1層目のカーネルサイズを3x3に縮小
CNN2層目を 1x1 conv & depthwise conv モジュールに変更
チャンネル数の縮小

見ての通りMobileNet v2にオマージュを効かせたネットワークとなっています。再度M5Cameraにデプロイしてみましょう。

目に見えて高速化しましたね。推論時間は約26msとなりました¹。

メモリフットプリント、バイナリサイズともまだかなり余裕があるため²、リアルタイム性を捨てより大規模な画像認識に挑戦するのも面白そうです。

まとめ

M5Cameraに搭載されたESP32で、CNNを使った数字認識を行った
Neural Network Librariesで学習したモデルをFreeRTOSに移植した
モデルの軽量化を行い、ほぼリアルタイムな推論を実現
もっと大規模なタスクも実行できそう

参考

M5Cameraのサンプルコードを実行してみる
http://mag.switch-science.com/2018/12/21/m5camera-test/

モデルアーキテクチャ観点からのDeep Neural Network高速化
https://www.slideshare.net/ren4yu/deep-neural-network-79382352

ESP8266で手書き文字を認識してみた
https://qiita.com/triwave33/items/6a4c02452dfffac4394a

環境

M5Camera (Bモデル)
Ubuntu 18.04.3 LTS
ESP-IDF v3.3
nnabla 1.0.10
nnabla-c-runtime 1.0.8

ここまで来ると推論以外の処理時間が気になってきますね。ESP32はデュアルコアのマイコンですが、実は今回のコードは大半の部分がシングルコアで動作しており、高速化の障壁になっています。上手くタスクが分散するよう書き直せば、非推論部もかなり高速化できるでしょう（未検証） ↩
厳密なベンチは取っていませんが、内臓フラッシュメモリは3MB、PSRAMは3.5MBほど余裕がありました。 ↩