Tensorflow 2.0 (Keras) を使ってニューラルネットワークをモデリングをし、学習済データを生成してSPRESENSEで動かしてみました。Tensorflowは、2.0以降のKerasになってからかなり分かりやすくなったのですが、マイコンへの移植は Tensorflow lite/micro のコードの分かり難さのせいで挫折している人は多いのではないでしょうか？

今回、KerasでAIをモデリングして学習済データを出力し、SPRESENSEで動かすことができましたので、その手順を紹介します。

Tensorflow の開発環境をセットアップする

Tensorflow は pip でインストールができます。対象バージョンは 2.8.0 です。バージョンを指定してインストールしてください。

tf_install.py

pip install tensorflow==2.8.0

Tensorflowでモデリングするときは、Jupyter notebook を使うと便利です。Jupyter notebookも、pipでインストールできます。検索するといろいろと使い方が出てきます。私もブログに使い方を紹介していますので、興味ある方は参照してください。

Jupyter notebook を使ってみた
https://makers-with-myson.blog.ss-blog.jp/2020-05-22

Tensorflow でモデリング

Tensorflow (Keras)をインポート

Tensorflow含め、必要なモジュールをインポートします。この時、Tensorflow のバージョンが'2.0.0'であることを確認してください。

tf_import.py

# TensorFlow and tf.keras
import tensorflow as tf
from tensorflow import keras

# Helper libraries
import numpy as np
import matplotlib.pyplot as plt
import random

print(tf.__version__)

データセットの準備

今回、認識させるのはお約束のMNISTです。データセットがあらかじめ準備されているのでこれを使います。デフォルトでは、60,000分のデータが学習用に10,000分がテスト用に割り振られます。

tf_mnist.py

mnist = keras.datasets.mnist
(org_train_images, org_train_labels), (org_test_images, org_test_labels) = mnist.load_data()
# 60,000 training data and 10,000 test data of 28x28 pixel images
print("train_images shape", org_train_images.shape)
print("train_labels shape", org_train_labels.shape)
print("test_images shape", org_test_images.shape)
print("test_labels shape", org_test_labels.shape)

データの正規化

入力は0.0-1.0の範囲とするので画像を正規化します。この時ラベルデータも浮動小数点に変換するのを忘れずに行いましょう。これもハマるポイントです。

tf_normalize.py

# Normalize the input image so that each pixel value is between 0 to 1.
train_images = org_train_images;
test_images = org_test_images;
train_images = train_images / 255.0
test_images = test_images / 255.0

train_labels = org_train_labels.astype("float32")
test_labels = org_test_labels.astype("float32")

print('Datasets are normalized')

ニューラルネットワークを設計し学習を実行する

ニューラルネットワークは標準的な畳み込みニューラルネットワークにしました。このあたりはKerasになって、とっても分かりやすくなりました。

tf_nnmodel.py

# Define the model architecture
model = keras.Sequential([
  keras.layers.InputLayer(input_shape=(28, 28)),
  keras.layers.Reshape(target_shape=(28, 28, 1)),
  keras.layers.Conv2D(filters=6, kernel_size=(5, 5), activation=tf.nn.relu),
  keras.layers.MaxPooling2D(pool_size=(2, 2)),
  keras.layers.Flatten(),
  keras.layers.Dense(32, activation=tf.nn.relu),
  keras.layers.Dense(10),
  keras.layers.Activation(tf.nn.softmax)
])

# Define how to train the model
model.compile(loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

# Train the digit classification model
model.fit(train_images, train_labels, batch_size=128, epochs=5, verbose=1)

学習済モデルを評価する

学習済モデルの評価します。評価する前にモデルの概要（サマリ）を出力しています。これは後程説明するTensorflow lite/micro で確保するメモリ量を見積る際の参考値になります。

tf_training.py

# Printout model summary
model.summary()

# Evaluate the model using all images in the test dataset.
test_loss, test_acc = model.evaluate(test_images, test_labels)
print('Test accuracy:', test_acc)

Tensorflow lite 用のモデルに変換する

学習済モデルが出来たので、Tensorflow lite用のモデルに変換しファイルに出力します。

tf_convert.py

# Convert Keras model to TF Lite format.
converter = tf.lite.TFLiteConverter.from_keras_model(model)

tflite_float_model = converter.convert()

# Show model size in KBs.
float_model_size = len(tflite_float_model) / 1024
print('Float model size = %dKBs.' % float_model_size)
# Save the model to disk
open('model.tflite', "wb").write(tflite_float_model)

Tensorflow lite micro用にファイルに出力

ここからさらに、Tensorflow lite micro用にヘッダーファイルに出力します。Tensorflowには、 xxd という変換用プログラムも用意されていますが、ファイルのエンコードがUTF-8でないなど不便が多いので、ここではPythonで出力します。

tf_model_header.py

import binascii

def convert_to_c_array(bytes) -> str:
  hexstr = binascii.hexlify(bytes).decode("UTF-8")
  hexstr = hexstr.upper()
  array = ["0x" + hexstr[i:i + 2] for i in range(0, len(hexstr), 2)]
  array = [array[i:i+10] for i in range(0, len(array), 10)]
  return ",\n  ".join([", ".join(e) for e in array])

tflite_binary = open("model.tflite", 'rb').read()
ascii_bytes = convert_to_c_array(tflite_binary)
header_file = "const unsigned char model_tflite[] = {\n  " + ascii_bytes + "\n};\nunsigned int model_tflite_len = " + str(len(tflite_binary)) + ";"
# print(c_file)
open("model_data.h", "w").write(header_file)

SPRESENSEに学習済モデルを組み込む

Spresense Tensorflow対応のArduino Packageをインストール

SPRESENSEにTensorflowの学習済モデルを組み込みます。出力したヘッダーファイルをスケッチのフォルダーにおいておきます。スケッチには、Tensorflow が組み込まれた　Spresense Arduino Board Package が必要です。

Spresense Arduino Package for Tensorflow
https://github.com/YoshinoTaro/spresense-arduino-tensorflow

Bitmap用のライブラリをインストールする

テスト用画像にビットマップファイルを使うので、Bmp用のライブラリも使用します。あらかじめインストールをしておきます。

BmpImage_ArduinoLib
https://github.com/YoshinoTaro/BmpImage_ArduinoLib

テスト用の画像を本体フラッシュに転送しておく

テスト用のデータは、本体のフラッシュに転送しました。フラッシュへの書き込みは、xmodem_writerを使います。次のURLから取得できます

Windows用
https://github.com/sonydevworld/spresense/blob/master/sdk/tools/windows/xmodem_writer.exe
Linux用
https://github.com/sonydevworld/spresense/blob/master/sdk/tools/linux/xmodem_writer
macOS用
https://github.com/sonydevworld/spresense/blob/master/sdk/tools/macos/xmodem_writer

次のコマンドでSpresense本体のフラッシュにデータを書き込めます。

xmodem.sh

$ xmodem_writer -c COM3 0009.bmp

拡張ボードをもっている人は、SDカードにコピーしてSDHCIライブラリを使った方が楽です。テスト用データをまとめたものを、ここに貼り付けておきますので使ってください。（pngを256色ビットマップに変換して使ってください）

スケッチの作成

Spresense のスケッチは次のようになりました。このスケッチはArduino用のBMPライブラリを使用しています。Githubからダウンロードして使ってください。

Tensorflow lite/micro では、tensor_arenaでTensorflow用のワークメモリを確保する必要がありますが、ここは実験的に求めるしかありません。一度大きめにメモリを確保したのちに、interpreter->arena_used_bytes() で実際に使うメモリ量を出力してサイズを調整します。

spresense_tf_mnist.cpp

#include "tensorflow/lite/micro/all_ops_resolver.h"
#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/micro/system_setup.h"
#include "tensorflow/lite/schema/schema_generated.h"

#include "model_data.h"

#define TEST_FILE "0009.bmp"

tflite::ErrorReporter* error_reporter = nullptr;
const tflite::Model* model = nullptr;
tflite::MicroInterpreter* interpreter = nullptr;
TfLiteTensor* input = nullptr;
TfLiteTensor* output = nullptr;
int inference_count = 0;

constexpr int kTensorArenaSize = 20000;
uint8_t tensor_arena[kTensorArenaSize];

#include <Flash.h>
#include <BmpImage.h>
BmpImage bmp;

void setup() {
  Serial.begin(115200);
  tflite::InitializeTarget();
  memset(tensor_arena, 0, kTensorArenaSize*sizeof(uint8_t));
  
  // Set up logging. 
  static tflite::MicroErrorReporter micro_error_reporter;
  error_reporter = &micro_error_reporter;
  
  // Map the model into a usable data structure..
  model = tflite::GetModel(model_tflite);
  if (model->version() != TFLITE_SCHEMA_VERSION) {
    Serial.println("Model provided is schema version " 
                  + String(model->version()) + " not equal "
                  + "to supported version "
                  + String(TFLITE_SCHEMA_VERSION));
    return;
  } else {
    Serial.println("Model version: " + String(model->version()));
  }
  
  // This pulls in all the operation implementations we need.
  static tflite::AllOpsResolver resolver;
  
  // Build an interpreter to run the model with.
  static tflite::MicroInterpreter static_interpreter(
      model, resolver, tensor_arena, kTensorArenaSize, error_reporter);
  interpreter = &static_interpreter;
  
  // Allocate memory from the tensor_arena for the model's tensors.
  TfLiteStatus allocate_status = interpreter->AllocateTensors();
  if (allocate_status != kTfLiteOk) {
    Serial.println("AllocateTensors() failed");
    return;
  } else {
    Serial.println("AllocateTensor() Success");
  }

  // Output the actual arena used bytes.
  size_t used_size = interpreter->arena_used_bytes();
  Serial.println("Area used bytes: " + String(used_size));

  input = interpreter->input(0);
  output = interpreter->output(0);

  Serial.println("Model input:");
  Serial.println("dims->size: " + String(input->dims->size));
  for (int n = 0; n < input->dims->size; ++n) {
    Serial.println("dims->data[" + String(n) + "]: " + String(input->dims->data[n]));
  }

  Serial.println("Model output:");
  Serial.println("dims->size: " + String(output->dims->size));
  for (int n = 0; n < output->dims->size; ++n) {
    Serial.println("dims->data[" + String(n) + "]: " + String(output->dims->data[n]));
  }
  
  /* read test data */
  File myFile = Flash.open(TEST_FILE);
  if (!myFile) { Serial.println(TEST_FILE " not found"); return; }

  Serial.println("Read " TEST_FILE);
  bmp.begin(myFile);
  BmpImage::BMP_IMAGE_PIX_FMT fmt = bmp.getPixFormat();
  if (fmt != BmpImage::BMP_IMAGE_GRAY8) {
    Serial.println("support format error");
    return;
  }

  int width = bmp.getWidth();
  int height = bmp.getHeight();
  Serial.println("width:  " + String(width));
  Serial.println("height: " + String(height));
  uint8_t* img = bmp.getImgBuff();
  // normalize the test image
  for (int i = 0; i < width*height; ++i) {
    input->data.f[i] = (float)(img[i]/255.0);
  }

  Serial.println("Do inference");

  TfLiteStatus invoke_status = interpreter->Invoke();
  if (invoke_status != kTfLiteOk) {
    Serial.println("Invoke failed");
    return;
  }

  /* output result */
  for (int n = 0; n < 10; ++n) {
    float value = output->data.f[n];
    Serial.println("[" + String(n) + "] " + String(value)); 
  }
}

void loop() { }

試してみた感想

今回作ったスケッチは、Tensorflow lite micro のサンプルコードよりもかなりすっきりしたと思います（自画自賛）。個人的な感想ですが、本家のサンプルはディレクトリ構成も酷くてとても読めたものではありません。組み込みで Tensorflow が広まらないは、あのサンプルのせいだと思っています。今回の取り組みでだいたい何をどうすればよいのか分かったので、これからSpresense/Tensorflow のサンプルを拡充したいと思います。