Tensorflow (Keras) の出力モデルを量子化しSPRESENSEで動かしてみた

Tensorflow は量子化モデルが使えるようになっています。通常、Tensorflowで出力するモデルは重みづけデータなどが浮動小数点のためマイコンシステムのメモリや計算資源を圧迫します。このモデルをより少ないビットで表現することを量子化モデルといいます。量子化を行うことによって、マイコンシステムで扱えるニューラルネットワークの幅を広げることができます。

前準備

SPRESENSE で Arduino IDE で Tensorflow lite/micro を使えるようにするには、専用のボードパッケージを使う必要があります。次の記事を参考にしてお使いの Arduino IDE にインストールをしてください。

また、SPRESENSEで Tensorflow のモデルを動かす場合は、バージョンは2.8.0である必要があります。（2022年4月3日現在）Anaconda3 などの Python 開発環境から、次のコマンドでインストールすることができます。

tf_install.py

$ pip install tensorflow==2.8.0

SPRESENSE で動作確認するためのテストデータに MNIST のデータを使います。次の画像をダウンロードして使うことができますが、画像フォーマットはPNGからBMP(256階調グレー画像)に変換してください。変換したBMP画像はSPRESENSEから読み込めるように、SDカードやSpresense本体の内蔵フラッシュに転送してください。

例題のニューラルネットワーク

今回は次のニューラルネットワークを使ってみます。MNISTの数字を認識する単純な畳み込みニューラルネットワークのモデルです。

keras_lite_model.py

import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt
import random
import binascii
print(tf.__version__)

# 60,000 training data and 10,000 test data of 28x28 pixel images
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Normalize the input images so that each pixel value is between 0 to 1.
train_images = (train_images / 255.0).astype("float32")
test_images = (test_images / 255.0).astype("float32")
train_labels = train_labels.astype("float32")
test_labels = test_labels.astype("float32")

# 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)

# Print out the 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)

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

# Show the 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)

出力されたモデルのサイズは112kBになりました。このモデルは浮動小数点の重みづけデータなどが格納されています。計算も浮動小数点で行っています。

output.sh

Test accuracy: 0.9794
Float model size = 112KBs.

量子化してサイズを小さくする

先ほどのニューラルネットワークを量子化します。Tensorflowには量子化のやり方はいくつかあるようですが、最も簡単なやり方である"ダイナミックレンジの量子化"を試してみました。これは中間層のモデルを浮動小数点32ビットから動的に8ビットの精度に量子化する手法です。量子化をするには、Tensorflow lite に変換するコードを次のように変更します。

quantize.py

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

def representative_dataset_gen():
   for i in range(100):
      input_image = tf.cast(test_images[i], tf.float32)
      input_image = tf.reshape(input_image, [1,28,28])
      yield ([input_image])

converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset_gen
# Convert the model
tflite_float_model = converter.convert()

この手法では、量子化の精度を上げるためにデータを与える必要があります。データを与えるための関数は別途定義する必要があり、この場合は、representative_dataset_gen関数がそれにあたります。

これを先ほどの畳み込みニューラルネットワークに組み込みます。

keras_lite_qmodel.py

import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt
import random
import binascii
print(tf.__version__)

# 60,000 training data and 10,000 test data of 28x28 pixel images
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Normalize the input image so that each pixel value is between 0 to 1.
train_images = train_images / 255.0
test_images = test_images / 255.0
train_labels = train_labels.astype("float32")
test_labels = test_labels.astype("float32")

# 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)

# Print out the 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)

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

# Quantize the model
def representative_dataset_gen():
   for i in range(100):
      input_image = tf.cast(test_images[i], tf.float32)
      input_image = tf.reshape(input_image, [1,28,28])
      yield ([input_image])

converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_float_model = converter.convert()

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

# evaluate the quantized model
# A helper function to evaluate the TF Lite model using "test" dataset.
def evaluate_model(interpreter):
  input_index = interpreter.get_input_details()[0]["index"]
  output_index = interpreter.get_output_details()[0]["index"]

  # Run predictions on every image in the "test" dataset.
  prediction_digits = []
  for test_image in test_images:
    # Pre-processing: add batch dimension and convert to float32 to match with the model's input data format.
    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)
    interpreter.set_tensor(input_index, test_image)

    # Run inference.
    interpreter.invoke()

    # Post-processing: remove batch dimension and find the digit with highest probability.
    output = interpreter.tensor(output_index)
    digit = np.argmax(output()[0])
    prediction_digits.append(digit)

  # Compare prediction results with ground truth labels to calculate accuracy.
  accurate_count = 0
  for index in range(len(prediction_digits)):
    if prediction_digits[index] == test_labels[index]:
      accurate_count += 1
  accuracy = accurate_count * 1.0 / len(prediction_digits)
  return accuracy

tflite_model_quant_file = "qmodel.tflite"
interpreter_quant = tf.lite.Interpreter(model_path=str(tflite_model_quant_file))
interpreter_quant.allocate_tensors()
print('Quantized model accuracy: ',evaluate_model(interpreter_quant))

これを出力すると、112kB あったモデルが 31kB まで圧縮できました。モデルのサイズは約1/4になっています。

output.sh

Float test accuracy: 0.976
Quantized model size = 31KBs.
Quantized model accuracy:  0.9765

量子化モデルをC言語ヘッダーとして出力

次のコードで量子化モデルをC言語ヘッダーに出力します。Tensorflow に付属している"xxd"コマンドでも変換できますが、Arduino IDEに読み込ませるには、"UTF-8"にしておく必要があるので注意してください。

c-header_out.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("qmodel.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("qmodel.h", "w").write(header_file)

SPRESENSE用のスケッチで動作確認

SPRESENSE用の次のスケッチで動作確認を行います。このスケッチはArduino用のBMPライブラリを使っています。次の Github からダウンロードして使ってください。

このテストスケッチは、SPRESENSEの内蔵フラッシュにテスト用画像が格納されていることを前提としています。拡張ボードを使っている人は、SDHCIライブラリに変更して使ってください。

spresense_tf_mnist.c

#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 "qmodel.h"
//#include "model.h"

#define TEST_FILE "0000.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");
  }

  size_t used_size = interpreter->arena_used_bytes();
  Serial.println("Area used bytes: " + String(used_size));
  input = interpreter->input(0);
  output = interpreter->output(0);

  /* check input */
  if (input->type != kTfLiteFloat32) {
    Serial.println("expected type is not float32");
    return;
  } else {
    Serial.println("input type is float32");
  }

  Serial.println("Model input:");
  Serial.println("dims->size: " + String(input->dims->size));
  Serial.println("dims->data[0]: " + String(input->dims->data[0]));
  Serial.println("dims->data[1]: " + String(input->dims->data[1]));
  Serial.println("dims->data[2]: " + String(input->dims->data[2]));
  Serial.println("input->type: " + String(input->type));

  Serial.println("Model output:");
  Serial.println("dims->size: " + String(output->dims->size));
  Serial.println("dims->data[0]: " + String(output->dims->data[0]));
  Serial.println("dims->data[1]: " + String(output->dims->data[1]));
  
  /* 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();

  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;
  }

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

void loop() { /* do nothing */ }

SPRESENSEで使用するメモリ量の比較

Tensorflowが必要とするメモリ量である tensor_arena の値を比較してみると、次のようになりました。もともとの浮動小数点のモデルに対し、1/3のサイズになっています。1/4となっていないのは入力と出力が浮動小数点となっているためと思われます。

モデル	モデルサイズ	必要メモリ量	認識精度
量子化前	112 kBytes	18192 Bytes	0.976
量子化後	31 kBytes	6096 Bytes	0.9765

メモリ削減の効果が確認でき、また認識精度もほぼ変わらないことが確認できました。これだけ小さくできると多少のサイズのニューラルネットワークモデルでもSPRESENSEで動かすことができそうです。

使ってみた感想

Tensorflowの量子化モデルを使ってみましたが、解説ドキュメントがちょっと分かり難いですね。でもこれが使えるようになるとマイコンで使えるニューラルネットワークの幅が広がるので、エッジAIのシステムを開発している人にとっては必須の技術と言えるでしょう。入力と出力も整数に置き換えることができる整数量子化（TFバージョンは2.3以上）を使えばメモリ量をもっと削減できるでしょう。その他、float16やint16の量子化など様々なバリエーションがあるので場面によって使い分けができるようになると便利だと思いました。