複数のNNライブラリを活用したAndroidアプリ内での深層ポーズ推定モデルの実行

Shintarou Okada

Engineer

本記事は、2020年インターンシップとして勤務した原悠輔さんによる寄稿です。

こんにちは。2020年夏季のPFNインターンに参加しました東大修士1年の原悠輔です。大学では、機械学習を利用したメディア処理を行う研究室に所属しています。

PFNインターンでは、深層ポーズ推定モデルを実行するAndroidアプリケーションの開発を行いました。ニューラルネットワーク（NN）モデル推論（、学習）用ライブラリであるTFLiteとLibTorchを協調させながら推論し、推定結果を元にUnity上でアバターを動かしています。TFLiteはTensorFlowのエッジデバイス向けツールで、LibTorchはPyTorchのC++向けインターフェースです。TFLiteとLibTorchの協調にはPFVMというPFN内製のNNモデル推論ライブラリを用いています。PFVMは既に公開しているchainer-compilerの後継にあたるものです。（経緯の詳細はこちら）

背景

スマホ上で深層学習モデルを実行することのメリットの1つは、実行用サーバの用意の必要がないことです。他にも、通信による遅延が生じないことや、オフラインで利用可能なこと、処理がデバイス内で完結しユーザにとってセキュリティ上の安心感があることが挙げられます。

一方で、スマホ上で利用可能なリソースはモデルの学習に用いるようなマシンと比べ貧弱です。プロセッサ性能やストレージサイズ、メモリサイズ、電力環境など、様々な面で劣ります。そのため、精度を落として計算を行ったり、GPU、DSP、TPUによるハードウェアアクセラレーションやモデル構造自体の最適化を施したりするなどの工夫が必要になります。

実装

今回実装した内容の概要を図2に示しました。NN推論ライブラリであるTFLiteとLibTorchを選択的に使いながら、GPU、CPUを用いて推論し、その結果を元にUnity上でアバターを動かすという構成になっています。

TFLiteに対して指定するオプションによっては思うようにハードウェアが使われなかったり、モデルの構造が特定の条件を満たす場合にTensorFlowグラフへの変換で失敗することが判明したりし、必要に応じてライブラリ内部の変更まで行うなど想定より試行錯誤を要しました。

ここからは各要素を順番に説明していきます。

TFLiteとLibTorch

TFLiteはTensorFlow Liteの略で、TensorFlowのモデルをスマートフォンやIoT機器などのエッジデバイスで実行するためにTensorFlow下で開発されているNNモデル推論ライブラリです。TFLiteは、エッジデバイスでの読み出しに適したファイル形式への変換や、CPU以外のハードウェアを使った実行、量子化（計算の精度を落とした実行）をサポートしています。デバイスのベンダーによる違いをTFLiteが吸収するので様々な端末で利用可能です。ただし、対応しているオペレータの種類が限られているため、複雑なモデルを動かすことは困難です。

一方、LibTorchはPyTorchのコア機能で、C++から利用することができます。LibTorchは、現状Android上でハードウェアアクセラレーションを利用できない一方で、PyTorchで利用可能なほとんどのオペレータに対応しています。

以上から、畳み込み層などの”重たい”オペレータはTFLiteで動かしつつ、TFLiteが対応していないオペレータを代わりにLibTorchに実行させることで複雑なモデルを高速に動かすことが可能になります。こうしてフレームワークを切り替えながら実行するために後述するPFVMというPFN内製のライブラリを使っています。

GPUアクセラレーションと量子化

今回作ったアプリケーションでは、TFLiteの機能でGPUアクセラレーションと実行時量子化をしています。これにより、単純にLibTorchのみを使ってCPUで動かした場合よりも全体で2倍高速に動作するようになりました。

TFLiteを使ってGPUアクセラレーションを有効にする方法は2通りあります。今回はTFLiteからNNAPIを経由してGPUを利用する方法をとりました。（NNAPIを経由せず直接GPUを利用する方法も試みましたが、今回のモデルが何らかの条件を満たしていなかったせいか十分な割合のオペレータがGPUで実行されず断念しました。）

NNAPIは、Android OSの提供するNNモデル推論用のAPIです。Android 8.1以上であればベンダーの違いを吸収して全てのデバイスで動作するとされています。GPUやTPU（Tensor Processing Unit）などのアクセラレータを、ベンダーの違いを意識せずに（ベンダーにより対応の程度の差はありますが）利用することが可能です。

量子化については、NNAPIの機能として実行時に単精度（32bit）浮動小数点数の計算を半精度（16bit）で行うオプションがあり、今回はこれを利用しました。

こうしたオプションを試すにあたっては、PFVMを介する前に簡単に実機上でTFLite単体での推論を行うため、TFLiteの公式で開発されているベンチマークツールも利用しました。これを用いると、TFLite形式のモデルをTFLiteを用いて実行した場合の推論時間をオペレータごとに計測することが可能です。CPU以外のハードウェアやNNAPIを用いて実行される部分については時間を計測できないという制限はあるものの、どの部分がどのハードウェアやNNAPIを利用して実行されるのかを知ることができ、例えば「想定より早く実行されない原因は十分にGPUを使えていないからだ」と特定できた場面もありました。

PFVM

図3 実行時にTFLiteとLibTorchが使い分けられている様子
（図中の時間は推論速度の限界ではありません。より軽量なモデルへの変更や、枝刈りをはじめとした工夫の余地が多数残されています。）

PFVMはchainer-compilerの後継にあたるPFN内製のNNモデル推論ライブラリです。ONNX形式で与えられたNNモデルに対してデプロイ前に

カーネル融合などの計算グラフの最適化を行う
TFLiteで実行可能な部分を検出し、どこをTFLiteで実行するか定める
定めた部分をonnx-tfコマンド、TFLiteコンバータを用いてTFLite形式のデータに変換する

ということを行い、デプロイ後の実行時には

定めた計画に従ってTFLiteとLibTorchを呼び出す

という役割を果たします。

図3は今回のポーズ推定モデルの推論時間の内訳を表しています。3つの断片がTFLiteによって実行されていることが分かります。

TFLiteで実行する部分をあまり細かく切り出すとTFLiteの呼び出しのオーバーヘッドの影響を無視できなくなるため、一定以上の大きさの部分グラフになるように切り出します。

これにより、ハードウェアアクセラレーションを使えるなどのTFLiteのメリットを活かしつつ、TFLiteが対応していないオペレータにLibTorchを使って対応することが可能になっています。今回はTFLiteを用いましたが、他にQualcomm社のSNPE（Snapdragon Neural Processing Engine）などのライブラリを選択することも可能です。

余談ですが、今回のモデルの用意にあたっては、PyTorchを用いて学習したものをpytorch-pfn-extras.onnxを使ってONNX形式にしています。pytorch-pfn-extrasはPyTorchを拡張してよりスムーズに研究開発を行うためのオープンソースのプロジェクトです。例えばpytorch-pfn-extras.onnxは、PyTorchとONNX形式の橋渡しをするAPIであるtorch.onnxの拡張であり、これを使うとモデルをONNX形式で出力する際にモデルだけでなく、入力と出力を含んだONNXのテストケースを作成できます。テストケースがあれば、推論が正しく行われるかのチェックが容易になります。

Unity

Unity上では、PFVMの呼び出しとアバターの操作を行います。Unityのスクリプトは一般にC#で書くため、C++で実装されているPFVMを呼び出すにはCのAPIに変換するラッパーを書いてプラグインとして読み込む必要がありました。

アバターの操作にあたっては、推定したポーズをもとに上半身の各関節の角度を算出し、アバターの各関節の角度に反映するという方法をとりました。

今後の課題

時間の都合上、今回のインターンで完遂できなかったタスクがいくつかあります。

ひとつは8bit整数量子化の適用です。

TFLiteには、モデルの重みをデプロイ前に8bit整数に量子化する機能があります。これにより、単精度（32bit浮動小数点数）で学習したモデルのサイズを1/4に減らし、アプリケーションのサイズを大幅に削減することができます。また、メモリ帯域の節約にも繋がり、メモリ帯域が律速になっている場合には実行の高速化を望むこともできます。残念ながら、今回は8bit整数に量子化したモデルを実機上で実行する段階で問題が発生し、期間中に終えることはできませんでした。

他には、モデルそのものの最適化があります。

モバイル環境に合わせてモデルを小さくして学習し直しても十分な精度が出る可能性があります。実際に、畳み込み層のチャンネル数を異なる割合で減らしたモデルを訓練し、実用上十分なチャンネル数を見積もるということをしました。畳み込み層に手を加えるのは、モデル全体に対して特に支配的な計算量・メモリ帯域を要する部分であることが多いからです。しかし残念ながら、時間の都合により実機にデプロイして精度や速度を確かめるまでは至りませんでした。

あるいは、学習中/学習後に枝刈りを行って推論結果にあまり影響しない部分の計算を省略できる可能性もあります。

モデルを小さくするにあたっては、NNAPIのプロファイリングも重要です。今回は、NNAPI呼び出し時のGPU/CPUの使われ方や各オペレータの所要時間を追うことができませんでしたが、これらの情報が得られれば狙いを定めたモデルサイズの調整が可能になります。

まとめ

今回のPFNインターンでは、PFVMを利用してTFLiteとLibTorchの各々の機能を活かしながら深層ポーズ推定モデルを実行するAndroidアプリケーションを実装しました。

C++に不慣れなために基本的なところで躓いたり、PFVMやonnx-tensorflow、TFLiteといったフレームワークの内部の処理を読んで変更する必要が生じたりと、一筋縄ではいきませんでしたが、メンターの岡田さん、渡辺さんや他の社員の方々（特にメンター並みに相談に乗ってくださった野崎さん）に助けていただきながらなんとか進めることができました。短い期間ではありましたが、学びの多い充実したインターンシップとなりました。ありがとうございました。