4ドルのマイコンで雑なAI推論をしようじゃないか

2024年に入ってからRaspberry Pi Pico向けのTensorFlow Lite for Microcontrollersが大幅アップデートされたので少し使ってみました。といっても、よくある画像や音声を判定する類の「なんかAIっぽい」ネタではなくもっと身近な題材で、ボタン操作をサンプリングして操作の状態をAIで推論します。大鉈で爪楊枝削ってる感じがしますが、4ドルのハードウェアで動かすTinyMLはTinyなネタでありたい。

プロジェクトの概要

このプロジェクトではRaspberry Pi Picoのボタン入力をシングルクリック、ダブルクリック、無操作の3つの状態に分類するニューラルネットワークモデルを構築します。TensorFlowで作成してトレーニングしたモデルをTensorFlow Lite microモデルに変換し、Picoで推論します。
つまり、モデルの作成と学習はPCで行い、そのモデルを組み込んだマイコンで推論します。

トレーニングデータの準備

マイコン向けの機械学習も類にもれず、データを用意し、前処理して、モデルに学習させ、評価して、実装に組み込むという流れになります。何はともあれまずはトレーニングデータ。データはRaspberry Pi Picoのボタン操作を32Hzで記録し、PythonのNumpy配列としてシリアルに出力したものを利用します。例としてシングルクリックのデータはこんな感じになりました:

single_click = np.array([
       [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

ダブルクリック:

double_click = np.array([
       [1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0],
       [1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
       [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0],
       [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]])

無操作はぜんぶ0:

nop = np.array([
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

それぞれ10件ずつの簡単なデータです。これだけでは流石に少ないので、データの配置をランダムに揺さぶってバリエーションを10倍にするデータ拡張を施します。これらトレーニングデータの前処理はホストPCのPythonで行いますが、詳細は割愛するのでGithubリポジトリを参照。

モデルの設計

入力は2値の時系列データ20個で、出力は [無操作, シングルクリック, ダブルクリック] の3状態なので、20個の特徴量を元に3クラスへ分類するニューラルネットワークモデルを定義します。無駄にDNN。かなり適当です。

from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.models import Sequential

model = Sequential([
    Dense(128, activation='relu', input_shape=(20,)),
    Dropout(0.5),
    Dense(64, activation='relu'),
    Dropout(0.5),
    Dense(3, activation='softmax')
])

データ拡張したトレーニングデータで学習したモデルをTensorFlow Liteモデルに変換し、さらにファイルシステムのないマイコン向けのモデルとしてCヘッダーファイルに変換します。

import tensorflow as tf

def convert_to_c_array(bytes_data):
    hex_array = [format(x, '#04x') for x in bytes_data]
    c_array = ''
    for i, hex_val in enumerate(hex_array):
        if i % 10 == 0 and i != 0:
            c_array += '\n'
        c_array += hex_val + ', '
    return c_array.strip().rstrip(',')


converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

print("Convert model to model.h file")
c_array = convert_to_c_array(tflite_model)
header_file_content = f"""
#ifndef MODEL_H
#define MODEL_H

const unsigned char model_tflite[] = {{
{c_array}
}};

const int model_tflite_len = {len(tflite_model)};

#endif  // MODEL_H
"""
with open('model.h', 'w') as f:
    f.write(header_file_content)

model.hはFlatBuffersでシリアライズしたTensorFlow LiteモデルをCの配列に変換したものです。

const unsigned char model_tflite[] = {
0x1c, 0x00, 0x00, 0x00, 0x54, 0x46, 0x4c, 0x33, 0x14, 0x00,
0x20, 0x00, 0x1c, 0x00, 0x18, 0x00, 0x14, 0x00, 0x10, 0x00,
0x0c, 0x00, 0x00, 0x00, 0x08, 0x00, 0x04, 0x00, 0x14, 0x00,
0x00, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x90, 0x00, 0x00, 0x00,
0xe8, 0x00, 0x00, 0x00, 0x08, 0xb0, 0x00, 0x00, 0x18, 0xb0,
...

あとは、Pico向けのTensorFlow Lite for Microcontrollers pico-tflmicroライブラリを使って、モデルを読み込んで推論するだけです。カンタン。

ボタン操作の推論

実際のセンサーデータの収集と推論はRaspberry Pi Picoで行います。pico-tflmicroライブラリでモデルを読み込み、サンプリングしたボタン操作のデータで推論します。推論結果はシリアルに逐次出力します。だいたいこんな出力:

Nop 1.00, Single click 0.00, Double click 0.00 -> Predicted label: Nop
Nop 1.00, Single click 0.00, Double click 0.00 -> Predicted label: Nop
Nop 1.00, Single click 0.00, Double click 0.00 -> Predicted label: Nop
Nop 0.00, Single click 1.00, Double click 0.00 -> Predicted label: Single click
Nop 0.00, Single click 1.00, Double click 0.00 -> Predicted label: Single click
Nop 0.00, Single click 0.00, Double click 1.00 -> Predicted label: Double click
Nop 0.00, Single click 0.00, Double click 1.00 -> Predicted label: Double click
Nop 0.00, Single click 0.00, Double click 1.00 -> Predicted label: Double click
Nop 0.00, Single click 0.91, Double click 0.09 -> Predicted label: Single click
Nop 0.00, Single click 0.99, Double click 0.01 -> Predicted label: Single click
Nop 0.00, Single click 0.01, Double click 0.99 -> Predicted label: Double click
Nop 0.00, Single click 0.00, Double click 1.00 -> Predicted label: Double click
Nop 0.00, Single click 0.03, Double click 0.97 -> Predicted label: Double click
Nop 0.00, Single click 0.07, Double click 0.93 -> Predicted label: Double click
Nop 0.00, Single click 0.01, Double click 0.99 -> Predicted label: Double click
Nop 0.00, Single click 0.27, Double click 0.73 -> Predicted label: Double click
Nop 0.00, Single click 1.00, Double click 0.00 -> Predicted label: Single click
Nop 1.00, Single click 0.00, Double click 0.00 -> Predicted label: Nop
Nop 1.00, Single click 0.00, Double click 0.00 -> Predicted label: Nop
Nop 1.00, Single click 0.00, Double click 0.00 -> Predicted label: Nop

ボタンの操作を元に推論する3状態それぞれの確率と、最も確率が高い状態を出力します。この実装は何も工夫せずにサンプリングしたデータをただ入力しているので、シングルクリックとダブルクリックが混ざってますが、適当な実装なのでその辺は実際に使う人が工夫してください。

通常の実装であればボタン操作などGPIOの入力状態の判定には状態マシンを利用すると思います。しかしながら 「プログラムを書かずにサンプルデータを10件ほど用意して、ニューラルネットに学習させ、活用する」 という"プログラミングモデル"は、やはりどこか甘美で魅力的です。もう少し上手く使えるようになりたいです。

ニューラルネットワークやTensorFlowに関しては世の中に情報が氾濫していますが、オンラインで読めるニューラルネットワークと深層学習が分かりやすいし手軽なので、とっかかりとしておすすめです。あとはChatGPT。