今日から始める Edge AI

はじめに

ディープラーニングに始まる AI ブームはまだまだ全然終わりそうもないですが、世の中の興味関心は良いモデルを作ることに加えて、モデルで推論することにも注目されてきているように思います。日本で MLOps という言葉が使われだしたのは昨年（2019 年）くらいで、ようやく機械学習を実用化する重要性が認識され始めて嬉しい限りです。機械学習を実用化するための重要な要素のひとつに、推論を実行する環境があります。Web サービスであればサーバサイド（クラウド）で実行することも可能ですが、ネットワーク帯域やリアルタイム性やセキュリティの問題で、デバイスやクライアントサイドで推論することもあります。コンピューティングリソースが豊富なサーバサイドに比べ、デバイスサイドではクラウドのようなスケールアウトはできませんし、電力消費も気にしなければなりません。サーバサイド推論とデバイスサイド推論でそれでも Edge AI を実用化するメリットや用途は多々あります。たとえば自動運転のようにネットワーク遅延を許容できない用途では、リアルタイム推論が可能な Edge AI が必須になります。スマホでも AR を用いたゲームでは、実世界を認識しつつ演出を表現するため、デバイスサイドでの演算が必要です。Edge AI はユーザの近くで便利なツールを提供するために重要な技術となってきているのです。

他方で Edge AI が広まっているかというと、まだまだ全然広まっていないように思えます。機械学習や AI 界隈で活動しているのに、Edge AI という単語を聞いたことがないというエンジニアもいるのではないでしょうか？　デバイスサイドで AI の推論を行う試み自体は、2016 年くらいには研究されてきていて、ディープラーニングのバイナリ化等で計算量を減らして推論を軽くする方法が提唱されています。モデルのバイナリ化は今でも有効な手段です。ディープラーニング専用のデバイス向けチップセット（Edge TPU 等）が開発されたり、Android/iOS に推論用チップセットが導入されたりと、Edge AI は密かに身近な存在になってきています。

前置きが長くなりましたが、今回の投稿では 2020 年冬時点の Edge AI の動向をまとめます。Edge AI は理論、ライブラリ、チップセット、実装含めて幅広い技術開発が進んでいます。全般的に整理することで、Edge AI の理解と実用化に寄与していきたいと思います。

Edge AI の全体像

Edge AI では推論をデバイスサイド（スマホや Web ブラウザやマシン）で実行します（スマホの中で学習する技術もありますが、今回は割愛します）。そのためにはデバイス側に学習済みモデルを配布し、モデルをロードして推論するランタイムが必要になります。スマホであれば TensorflowLite や PyTorch Mobile や MediaPipe がその技術になります。加えてディープラーニングのテンソル演算を効率的に実行する演算器も必要になります。サーバサイドであれば GPU を使いますが、デバイスサイドでは GPU か、テンソル演算に特化させたチップセットや FPGA を使います。デバイスは一般的にサーバサイドに比べて計算リソースが貧弱なため、軽量なモデルの活用や、モデル自体の軽量化（計算量削減）が必要になります。MobileNet に代表される軽量モデルや、バイナリ化や枝刈り（Pruning）に代表される軽量化が活用されます。配布するモデルの管理も必要です。モデルの再学習やアルゴリズムの変更によって新しいモデルを使う必要がある場合、サーバサイドであればサーバを停止して（または停止せずに）モデルを入れ替えることが可能です。しかし Edge AI の場合、モデルがデバイス側に配布されてしまっているため、モデルの更新にはデバイスでのモデルダウンロードが必要になります。配布システムには ML Kit（旧 Firebase ML Kit）や AWS Greengrass を活用することができます。

これらの技術について、以下では具体例含めて解説していきます。

スマホの Edge AI

スマホは Edge AI の重要な使い所の一つです。スマホではカメラやテキスト入力、音声入力で Edge AI が活用されます。スマホで AI を動かすためのライブラリは、汎用的なライブラリでは Tensorflow Lite、PyTorch Mobile、MediaPipe が用意されています。

Tensorflow Lite

Tensorflow Lite は Tensorflow をベースにしたデバイス推論用のランタイムです。Android、iOS、マイコン で使用可能です。Tensorflow で学習したモデルを TFLite Converter で変換して Tensorflow Lite 用のモデルファイル（.tflite）を生成することができます。量子化や枝刈りにも対応しており、OS に拘らずにスマホで推論する場合は最有力な選択肢になるライブラリです。加えて豊富なサンプルモデルやアプリが提供されており、自ら学習せずともある程度動作可能なモデルを使える点は魅力的です。
サンプルモデル
サンプルアプリ

TFLite Converter は Python API とコマンドラインツールが提要されています。公式には Python API の使用が推奨されています。TFLite Converter では Tensorflow や tf.keras の学習済みモデルを TFLite 用に変換します。書き方は以下のようになります。

import tensorflow as tf

# saved_modelを変換
converter = tf.lite.TFLiteConverter.from_saved_model(export_dir)
tflite_model = converter.convert()
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

# kerasモデルを変換
keras_model = tf.keras.models.load_model(filepath)
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
keras_tflite_model = converter.convert()
with open('keras_model.tflite', 'wb') as f:
    f.write(keras_tflite_model)

変換されたモデルは Android や iOS で tflite ライブラリでロードして推論に使うことができます。Android であれば以下のようなコード例になります。

class TFLiteActivity : AppCompatActivity() {
    /*モデルをロード*/
    private fun initializeTFLite(device: String = "NNAPI", numThreads: Int = 4) {
        val delegate = when (device) {
            "NNAPI" -> NnApiDelegate()
            "GPU" -> GpuDelegate()
            "CPU" -> "" }
        if (delegate != "") tfliteOptions.addDelegate(delegate)

        tfliteOptions.setNumThreads(numThreads)
        tfliteModel = FileUtil.loadMappedFile(this, tflite_model_path)
        tfliteInterpreter = Interpreter(tfliteModel, tfliteOptions)
        inputImageBuffer = TensorImage(tfliteInterpreter.getInputTensor(0).dataType())
        outputProbabilityBuffer = TensorBuffer.createFixedSize(
            tfliteInterpreter.getOutputTensor(0).shape(),
            tfliteInterpreter.getInputTensor(0).dataType())

        probabilityProcessor = TensorProcessor
            .Builder()
            .add(NormalizeOp(0.0f, 1.0f))
            .build()
    }

    /*推論*/
    @WorkerThread
    override fun analyzeImage(image: ImageProxy, rotationDegrees: Int): Map<String, Float> {
        val bitmap = Utils.imageToBitmap(image)
        val cropSize = Math.min(bitmap.width, bitmap.height)
        inputImageBuffer.load(bitmap)
        val inputImage = ImageProcessor
            .Builder()
            .add(ResizeWithCropOrPadOp(cropSize, cropSize))
            .add(ResizeOp(224, 224, ResizeMethod.NEAREST_NEIGHBOR))
            .add(NormalizeOp(127.5f, 127.5f))
            .build()
            .process(inputImageBuffer)

        tfliteInterpreter.run(inputImage!!.buffer, outputProbabilityBuffer.buffer.rewind())
        val labeledProbability: Map<String, Float> = TensorLabel(
            labelsList, probabilityProcessor.process(outputProbabilityBuffer)
        ).mapWithFloatValue
        return labeledProbability
    }
}

コード量も少なく Tensorflow Lite を使用できるため、Edge AI を始めるにはちょうど良いライブラリになっていると思います。

演算については Tensorflow の一部の演算子は Tensorflow Lite ではサポートされていません。
https://www.tensorflow.org/lite/guide/ops_compatibility

PyTorch Mobile

PyTorch Mobile は PyTorch をベースにしたスマホ推論用のランタイムです。これも Android、iOS をサポートしています。PyTorch Mobile は Pytorch のモデルを torchscript で jit コンパイルして.pt ファイルとして使う様になっています。2019 年にリリースされて以降、あまりアップデートがないのが不安でしたが、2020 年の Torch Developer Day で Android の NNAPI、iOS の Metal API のサポートを発表しました。それまで CPU のみで推論可能だった PyTorch Mobile が、GPU や専用プロセッサへのデリゲートが可能になりました。モデル学習の最有力が Pytorch である現場では、Pytorch で学習したモデルをスマホにインストールしてパフォーマンスが出るなら使いたい次第です。

PyTorch ではモデルの変換を torch.jit.trace で実行します。

import torch
import torchvision

model = torchvision.models.resnet18(pretrained=True)
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")

変換されたモデルは Android や iOS で tflite ライブラリでロードして推論に使うことができます。Android であれば以下のようなコード例になります。

class PyTorchActivity : AppCompatActivity() {
    /*モデルをロード*/
    private fun initializePyTorch() {
        val pytorchModule = Module.load(Utils.assetFilePath(
            this,
            pytorch_mobile_model_path))
        val mInputTensorBuffer = Tensor.allocateFloatBuffer(3 * 224 * 224)
        val mInputTensor = Tensor.fromBlob(
            mInputTensorBuffer,
            longArrayOf(1, 3, 224L, 224L)
        )
    }

    /*推論*/
    @WorkerThread
    override fun analyzeImage(image: ImageProxy, rotationDegrees: Int): Map<String, Float> {
        TensorImageUtils.imageYUV420CenterCropToFloatBuffer(
            image.image,
            rotationDegrees,
            224,
            224,
            TensorImageUtils.TORCHVISION_NORM_MEAN_RGB,
            TensorImageUtils.TORCHVISION_NORM_STD_RGB,
            mInputTensorBuffer,
            0
        )
        val outputModule = pytorchModule.forward(IValue.from(mInputTensor)).toTensor()
        val scores = outputModule.dataAsFloatArray
        val labeledProbability: MutableMap<String, Float> = mutableMapOf()
        for (i in 0 until labelsList.size - 1) {
            labeledProbability[labelsList[i + 1]] = score[i]
        }
        return labeledProbability
    }
}

Android の PyTorch Mobile はSoLoaderという Facebook 製のネイティブコードローダーで実行されています。

MediaPipe

続いて MediaPipe です。MediaPipe は Tensorflow Lite や PyTorch Mobile と違い、特定のディープラーニングライブラリから派生したものではありません。Google が開発しているスマホや Web ブラウザ向けのクロスプラットフォームなデータストリーミングエンジンで、スマホのカメラやマイクをインターフェイスとして入力された生データのデコーディングから前処理、推論、後処理までを一連のストリーミングとしてカバーするランタイムになっています。使えるモデルは Tensorflow Lite に制限されますが、推論時に課題となるデータ入力、前処理、後処理と推論の隔たりを統合しようとするライブラリで、使い勝手さえ良ければ（Bazel に依存しなければ）最高のランタイムになると思います。

スマホアプリに推論を組み込む UI は多々ありますが、MediaPipe ではその活用例も豊富に公開しています。
サンプル

MediaPipe では以下のような計算グラフを設計し、この計算グラフを bazel で変換して使用します。変換結果は Android であれば.aar ファイル、iOS であれば.ipa ファイルになります。

# MediaPipe graph that performs face mesh with TensorFlow Lite on GPU.

# GPU buffer. (GpuBuffer)
input_stream: "input_video"

# Output image with rendered results. (GpuBuffer)
output_stream: "output_video"
# Detected faces. (std::vector<Detection>)
output_stream: "face_detections"

# Throttles the images flowing downstream for flow control. It passes through
# the very first incoming image unaltered, and waits for downstream nodes
# (calculators and subgraphs) in the graph to finish their tasks before it
# passes through another image. All images that come in while waiting are
# dropped, limiting the number of in-flight images in most part of the graph to
# 1. This prevents the downstream nodes from queuing up incoming images and data
# excessively, which leads to increased latency and memory usage, unwanted in
# real-time mobile applications. It also eliminates unnecessarily computation,
# e.g., the output produced by a node may get dropped downstream if the
# subsequent nodes are still busy processing previous inputs.
node {
  calculator: "FlowLimiterCalculator"
  input_stream: "input_video"
  input_stream: "FINISHED:output_video"
  input_stream_info: {
    tag_index: "FINISHED"
    back_edge: true
  }
  output_stream: "throttled_input_video"
}

# Subgraph that detects faces.
node {
  calculator: "FaceDetectionFrontGpu"
  input_stream: "IMAGE:throttled_input_video"
  output_stream: "DETECTIONS:face_detections"
}

# Converts the detections to drawing primitives for annotation overlay.
node {
  calculator: "DetectionsToRenderDataCalculator"
  input_stream: "DETECTIONS:face_detections"
  output_stream: "RENDER_DATA:render_data"
  node_options: {
    [type.googleapis.com/mediapipe.DetectionsToRenderDataCalculatorOptions] {
      thickness: 4.0
      color { r: 255 g: 0 b: 0 }
    }
  }
}

# Draws annotations and overlays them on top of the input images.
node {
  calculator: "AnnotationOverlayCalculator"
  input_stream: "IMAGE_GPU:throttled_input_video"
  input_stream: "render_data"
  output_stream: "IMAGE_GPU:output_video"
}

Visualizerというツールで pbtxt で表現したグラフを可視化することが可能です。

上記は顔認識のグラフになっています。入力画像に対して顔を検知（Detection）し、検知した位置にオーバレイ（DetectionsToRender→AnnotationOverlay）を被せて画面に出力する、という一連のプロセスを表現しています。MediaPipe の良いところは、Edge AI を推論だけで終わらせず、入力 → 前処理 → 推論 → 後処理 → 出力までの一連の流れを組み込める点です。

ここまでスマホのランタイムについて書いてきましたが、OS（iOS、Android）側の動きを書いていきます。

iOS

まず iOS では各種 ML モデルをアプリに統合するためのフレームワークとして CoreML を提供しています。CoreML は Tensorflow Lite のようなスマホ向けのモデルだけでなく、ONNX や Pytorch、XGBoost、Scikit-learn といった多様なライブラリのモデルの推論をサポートしたツールです。モデルを coremltools で CoreML フォーマットに変換することで、ネイティブアプリにロードして推論を行うことができるようになっています。演算は自動的に（？）ニューラルエンジンという iPhone 特有の演算器にデリゲートされるようになっており、高速な推論と低消費電力を実現しています。更に一部のモデル（ニューラルネットワークの分類や最近傍）をデバイス側で更新（fine tuning）することができるようになっています。iOS で機械学習を使いたい場合は CoreML で始めるのが最も良い選択肢でしょう。

Android

Android では推論用の API として NNAPI（Android Neural Networks API）を提供しています。NNAPI は機械学習の処理に特化した Android C API で、Android8.1（API レベル 27）以降で使用可能です。NNAPI ではデバイスのハードウェア性能や負荷状況に応じてニューラルネットワークの処理を適切なデバイス（GPU、DSP、専用プロセッサ）にデリゲート（委譲）します。デリゲートできない処理は CPU で実行する仕組みです。iOS と違って多様なデバイスが販売される Android 特有の戦略と言えます。

Web ブラウザ

Tensorflow.js

Web ブラウザでも AI で学習や推論することが可能です。Web ブラウザで AI を動かす主な言語は Javascript で、有力なライブラリはTensorflow.jsになります。学習も推論も Web ブラウザ（端末）で実行することになるため、スマホにおける Edge AI のように端末のリソースをフル活用することはできません。しかし Web ブラウザが有する演算基盤としてWebGLやWASMを有効活用することで高速化を実現します。
推論モデルにはtf.keras モデルやsaved modelを Tensorflow.js 用に変換して使用することが可能です。モデルはニューラルネットワーク構造や Weight 含めて JSON 形式に変換され、クライアント端末にダウンロードされて使用されます。

変換にはコマンドラインツールのtensorflowjs_converterを使います。
変換例は以下のとおりです。

# saved modelを変換
tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_node_names='MobilenetV1/Predictions/Reshape_1' \
    --saved_model_tags=serve \
    /mobilenet/saved_model \
    /mobilenet/web_model

# Kerasモデルを変換
tensorflowjs_converter \
    --input_format keras \
    path/to/my_model.h5 \
    path/to/tfjs_target_dir

サンプルデモにあるとおり、Web ブラウザでも多様な体験を作ることが可能です。
Tensorflow.js のコード例は以下になります。Javascript として HTML に組み込むことが可能です。（参考元）

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.0.1"> </script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet@1.0.0"> </script>

<img id="img" src="cat.jpg"></img>

<script>
  const img = document.getElementById('img');
  // Load the model.
  mobilenet.load().then(model => {
    // Classify the image.
    model.classify(img).then(predictions => {
      console.log('Predictions: ');
      console.log(predictions);
    });
  });
</script>

まはた.js ファイルとして記述することも可能です。（参考元）

import * as tf from "@tensorflow/tfjs";

import { IMAGENET_CLASSES } from "./imagenet_classes";

const MOBILENET_MODEL_PATH =
  "https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json";

const IMAGE_SIZE = 224;
const TOPK_PREDICTIONS = 10;

let mobilenet;
const mobilenetDemo = async () => {
  status("Loading model...");

  mobilenet = await tf.loadLayersModel(MOBILENET_MODEL_PATH);
  mobilenet.predict(tf.zeros([1, IMAGE_SIZE, IMAGE_SIZE, 3])).dispose();

  status("");

  const catElement = document.getElementById("cat");
  if (catElement.complete && catElement.naturalHeight !== 0) {
    predict(catElement);
    catElement.style.display = "";
  } else {
    catElement.onload = () => {
      predict(catElement);
      catElement.style.display = "";
    };
  }
};

async function predict(imgElement) {
  status("Predicting...");
  const logits = tf.tidy(() => {
    const img = tf.browser.fromPixels(imgElement).toFloat();

    const offset = tf.scalar(127.5);

    const normalized = img.sub(offset).div(offset);

    const batched = normalized.reshape([1, IMAGE_SIZE, IMAGE_SIZE, 3]);

    return mobilenet.predict(batched);
  });
}

ml5.js

Tensorflow.js の使い方はスクリプトタグから Tensorflow.js をそのまま使用する方法と、ml5.jsから使用する方法の 2 通りがあります。ml5.js は Web 向けに機械学習をより使いやすくするためのライブラリです。ml5.js はp5.jsというデザイナー向けの Javascript と親和性高いコーディングになっています。ml5.js では画像や言語、音声といったメディアで頻繁に使われる分類や変換の API を用意しており、より感覚的に Javascript で AI を使えるようになっています。

ml5.js による推論は以下のようなコードになります。(参考元)

let classifier;

let img;

function preload() {
  classifier = ml5.imageClassifier("MobileNet");
  img = loadImage("images/bird.png");
}

function setup() {
  createCanvas(400, 400);
  classifier.classify(img, gotResult);
  image(img, 0, 0);
}

function gotResult(error, results) {
  if (error) {
    console.error(error);
  } else {
    console.log(results);
    createDiv(`Label: ${results[0].label}`);
    createDiv(`Confidence: ${nf(results[0].confidence, 0, 2)}`);
  }
}

モデルと軽量化

デバイス（スマホ、家電、自動車、機械等々）で推論するためにはモデルを軽量化する工夫が必要です。デバイスは一般的にサーバほど計算リソースや電力が豊富ではないため、計算量やモデルの容量を削減することで、デバイスの継続や応答性能を確保します。軽量化の工夫としては量子化(quantization)、枝刈り(pruning)、共有(sharing)、蒸留(distillation)が有名です。

量子化では重みのビット幅を軽量なものに削減します。たとえばモデルを作成する時点では float32 で学習し、推論前にモデルを float16 や int8 や 1bit(boolean)に変換することでモデルの容量や計算量を削減することができます。量子化によってニューラルネットワークの演算が学習時と推論時で多少ずれるため、計算効率化とトレードオフで精度が劣化する可能性があります。プロダクトとして高速化と精度劣化の妥協点を求めることが重要です。

枝刈りはニューラルネットワークのノードから重要性の低いものを削除する手法です。ネットワークの一部を削除することで、モデルを軽量化し、計算量や容量を削減します。量子化同様、効率化とトレードオフで精度が劣化する可能性があります。枝刈りの対象となるノードは一般的に重みが低いノードを言われています。枝刈りでは学習と枝刈りを繰り返し行うことで精度の維持と計算量の削減を目指します。

共有では重みを複数のノード間で共有します。重みを共有することでモデルの容量を削減することができます。他方でモデルの計算量は削減されないため、高速化が見込めるものではありません。

蒸留では学習方法を工夫することで軽量モデルの精度を改善する手法です。蒸留では最初に精度の良い大容量モデルで学習を行い、大容量モデルの推論結果を目的変数とします。大容量モデルが学習した正解データ（Hard target）と推論結果（Soft target）を用いて、軽量モデルを大容量モデルに近づくように学習させる手法が蒸留です。大容量モデルの推論結果は各ラベルに対する確率で表現されるため、軽量モデルはラベルの確率から各データのラベルに対する類似性を学習します。例えばあるネコの画像があり、Soft target としてネコ 60%、イヌ 30%、ウサギ 10%という分布になっていた場合、そのネコ画像がネコ:イヌ:ウサギ=6:3:1 という特徴を持っていると学習するというものです。

サービス

Edge AI を支えるサービスで有名なのは Google MLKit（旧 Firebase MLKit）です。MLKit は Android/iOS 向けにモデルの配布や推論、学習、ログ収集といった Edge AI に必要なサービスを提供する SDK です。MLKit から学習済みモデルや独自モデルを Android/iOS にダウンロードし、MLKit の SDK が提供するライブラリで、アプリ内で容易に推論を実行することが可能です。スマホ端末へのモデルの配布という Edge AI の課題を解決し、推論ラッパーまで提供されている便利なサービスになります。現在サポートされているモデルは Tensorflow Lite に限られます。

MLKit の便利な点は汎用的に使えるモデルを標準提供している点です。自らモデルを用意しなくても、MLKit の提供する API からモデルを取得して活用することができます。たとえば物体検知を使いたい場合、SSD や YOLO のモデルを用意しなくても、以下のコードで実行することが可能です。

private class ObjectDetection : ImageAnalysis.Analyzer {
    val options = FirebaseVisionObjectDetectorOptions.Builder()
        .setDetectorMode(FirebaseVisionObjectDetectorOptions.STREAM_MODE)
        .enableClassification()
        .build()
    val objectDetector = FirebaseVision.getInstance().getOnDeviceObjectDetector(options)

    private fun degreesToFirebaseRotation(degrees: Int): Int = when(degrees) {
        0 -> FirebaseVisionImageMetadata.ROTATION_0
        90 -> FirebaseVisionImageMetadata.ROTATION_90
        180 -> FirebaseVisionImageMetadata.ROTATION_180
        270 -> FirebaseVisionImageMetadata.ROTATION_270
        else -> throw Exception("Rotation must be 0, 90, 180, or 270.")
    }

    override fun analyze(imageProxy: ImageProxy?, degrees: Int) {
        val mediaImage = imageProxy?.image
        val imageRotation = degreesToFirebaseRotation(degrees)
        if (mediaImage != null) {
            val image = FirebaseVisionImage.fromMediaImage(mediaImage, imageRotation)
            objectDetector.processImage(image)
                    .addOnSuccessListener { detectedObjects ->
                        for (obj in detectedObjects) {
                            val id = obj.trackingId
                            val bounds = obj.boundingBox
                            val category = obj.classificationCategory
                            val confidence = obj.classificationConfidence
                            // Do Something
                        }
                    }
                    .addOnFailureListener { e ->
                        // Do Something
                    }
        }
    }
}

まとめ

Edge AI の全体像と使い方を図とサンプルコードで説明しました。Edge AI はスマホだけでなく、自動車や機械、家電等々、様々な場面で活躍する可能性を秘めたテクノロジーです。技術的にも今後成長していく領域なので、ぜひ広まってほしいと思っています。