【エッジAI】TensorFlow Liteを活用したRaspberry Piでのリアルタイム物体検出の実装

プログラミング

近年のAI開発では、クラウドサーバーで重い推論を行わせるだけでなく、カメラやセンサーが搭載された「現場のデバイス(エッジデバイス)」上で直接AIモデルを動かす「エッジAI」の重要性が急速に高まっています。遅延(レイテンシ)の短縮、プライバシー保護、オフライン環境での稼働など、クラウドAPIに頼らない分散型アーキテクチャはIoTの要です。

本記事では、ディープラーニングモデルを軽量化するフレームワーク「TensorFlow Lite (TFLite)」を活用し、定番のシングルボードコンピュータであるRaspberry Pi上でリアルタイムに高精度な物体検出を行うシステムの構築手法を解説します。

1. エッジデバイスへの学習済みモデルの最適化とデプロイ

PC上でKerasやTensorFlowを使って学習した巨大なニューラルネットワーク(CNN等)は、数百MB以上のサイズになりメモリ不足でRaspberry Piでは高速に動作しません。
これを解決するのが「量子化(Quantization)」と呼ばれる技術です。通常32ビットの浮動小数点(float32)で保持しているモデルの重みパラメータを、8ビット整数(int8)等に圧縮することで、モデルのサイズを4分の1に軽量化し、CPU/DSP上での推論速度を劇的に跳ね上げます。

import tensorflow as tf
# 通常のSavedModelやKerasモデルをロード
converter = tf.lite.TFLiteConverter.from_saved_model('saved_model_dir')
# フル整数(INT8)量子化の最適化フラグを設定
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 代表データセット(Representative Dataset)のジェネレータを提供(INT8量子化に必須)
def representative_data_gen():
    for input_value in tf.data.Dataset.from_tensor_slices(train_images).batch(1).take(100):
        yield [input_value]
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
# TFLiteモデルへの変換と保存
tflite_quant_model = converter.convert()
with open('quantized_model.tflite', 'wb') as f:
    f.write(tflite_quant_model)

2. Raspberry Pi上での推論ループの実装

モデルの軽量化が済んだら、Raspberry PiにTensorFlow Lite Runtime(フルのTensorFlowではなく、推論に特化した超軽量ライブラリ)とOpenCVをインストールし、カメラ映像の各フレームに対して推論を実行していくループを構築します。

import cv2
import numpy as np
import tflite_runtime.interpreter as tflite
# TFLite Interpreterの初期化
interpreter = tflite.Interpreter(model_path='quantized_model.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# カメラストリームのオープン
cap = cv2.VideoCapture(0)
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
        
    # 前処理:モデルの入力サイズ(例えば300x300)にリサイズし、次元を追加
    img_resized = cv2.resize(frame, (300, 300))
    input_data = np.expand_dims(img_resized, axis=0)
    # 推論の実行
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    # 出力テンソルの取得(バウンディングボックス位置、クラスID、スコア)
    boxes = interpreter.get_tensor(output_details[0]['index'])[0]
    classes = interpreter.get_tensor(output_details[1]['index'])[0]
    scores = interpreter.get_tensor(output_details[2]['index'])[0]
    # スコアが閾値を超えるものを描画
    for i in range(len(scores)):
        if scores[i] > 0.5:
            # 正規化された座標を元の画像サイズにスケールアップしてバウンディングボックスを描画...
            pass
    cv2.imshow('Object Detection on Edge', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

エッジAI開発のおすすめハードウェア機材

PC上のPython環境と違って、エッジAIの構築には実際にLinuxが動く物理的なSoCボードが必要です。
まずは定番中の定番である「Raspberry Pi 4 (メモリ4GBまたは8GBモデル)」を用意するのが最もコミュニティの情報が多く、トラブルシューティングしやすいです。

さらに、Raspberry PiのCPUだけでは推論FPSが実用レベルに届かない場合は、Googleが開発したエッジTPU(Tensor Processing Unit)である「Coral USB Accelerator」をType-Cで外付けすることで、推論の計算のみを劇的に高速化させることができます。TFLiteモデルをエッジTPU用にコンパイルするだけで、秒間30FPSクラスのリアルタイム処理が可能になります。

コメント

タイトルとURLをコピーしました