物体検出器 EfficientDet をブラウザで高速推論

R&D チームの奥村(@izariuo440)です。EfficientDet がブラウザで動いているのを見たことがなかったので、やってみました。以下はブラウザで実行中の様子1です。

結果として、EfficientDet-D0 で 256x256 の解像度の画像を TensorFlow Lite で推論させると、ネイティブで 20 FPS、WebAssembly SIMD support ありの Chrome で 7 FPS 程度の速度性能を確認できました。

Case XNNPACK XNNPACK スレッド数 FPS
A OFF N/A 4
B ON 1 15
C ON 4 20

測定環境は MacBook Pro (Retina, 15-inch, Mid 2015) の 2.5 GHz です。

背景と動機

2020/01/30、V8 開発ブログFast, parallel applications with WebAssembly SIMD という記事が投稿されました。2020/02/05、TensorFlow Lite に XNNPACK取り込まれました。2020/03/11、The TensorFlow BlogIntroducing the WebAssembly backend for TensorFlow.js、2020/03/16 に Higher accuracy on vision models with EfficientNet-Lite という記事が投稿されました。

これらのことから、GPU や専用ハードウェアではなく、CPU での推論のさらなる実用化に向けてコード最適化とモデル最適化の両面が進捗しており、かつそれをブラウザにも持ち込もうという動向を感じました。手元で少し試したところ、XNNPACK を有効にして TFLite で推論すると、CPU でもこんなに高速に推論できるのか、と驚きました。さらに WebAssembly SIMD ではブラウザでもこんなに高速に推論できるうのか、と驚きました。

エッジ推論においては NVIDIA の GPU 以外にも群雄割拠状態が続いています。モバイルでは iOS では CoreML、Android では NN API、他にも Qualcomm の Snapdragon Neural Processing Engine、最近では ARM の Arm NN など多種多様です。昨年度に登場した EdgeTPU は TensorFlow Lite ベースです。TensorFlow Lite はかなりのポータビリティがあり、NVIDIA CUDA や RasPi の Video Core を除けば多種多少なハードウェアに対応しています。TensorFlow Lite は CPU での推論が可能で、Delegate によってさまざまな環境に対応できるようになっており、現状の実装でも一部のハードウェアを除けばポータビリティの大きい現実的な選択肢です。人気のある深層学習フレームワークのひとつである PyTorch も、PyTorch Modile でモバイルでの推論のサポートに向かっています。

というわけで、こうした動向に備えて TensorFlow Lite で XNNPACK を有効にした場合の性能を試したりしていたときに、物体検出器として注目していた EfficientDet の公式実装が 2020/03/14 以降 GitHub で公開され始めたので、少しずつ取り組んでいかねば、と感じていました。私の実験用レポジトリのコミットログを見ると、2020/04/02 に WebAssembly + SIMD で EfficientDet が動き始めたようです。

あと、EfficientDet をブラウザで動かしているのを見たことがなかったので、やってみたかったのです。

周辺知識

EfficentDet

EfficientDet は、物体検出器2のひとつです。2019/11/20 に arXiv に論文が投稿され、2020/04/15 に Google AI Blog に記事が投稿されました。すでに多くの解説記事があるので詳細は省きますが、EfficientNet を物体検出タスクに持ち込み、速度と精度のトレードオフを調整しやすい RetinaNet 系の物体検出器、として位置づけられると思います。D0 から D7 までの8種類のモデルがあり、用途に応じた選択が可能というのも強いです。モバイル向けに EfficientNet-lite を backbone にすればさらに高速・軽量になりそうですし、2019 年度に台頭してきた CenterNet、Objects as Points などのキーポイントベースのアプローチを組み合わせればもっと性能が上がるかもしれない、という印象を持っています。

EfficientDet の精度

精度は公式 TensorFlow 実装 によると、COCO データセットの test-dev2017 において mAP が 33.8 〜 53.4 と驚異的ですし、mAP が 33.8 の EfficientDet-D0 は FP32 の訓練済みモデルのサイズが 16 MB 程度という軽量さも驚異的だと思います。また、EfficientDet-D0 でも入力解像度を大きくすることで、かなり小さい物体を検出できます。

EfficientNet の速度

速度は公式 TensorFlow 実装 によると、Tesla V100 で TensorRT なしで EfficientDet-D0 が 97 FPS、EfficientDet-D7 が 8 FPS のようです。EfficientDet の少し古い版の論文では、CPU の速度についても記載があったのですが、最近なくなってしまったようです。

RefineDet との比較

私の好きな RefineDet512 (VGG-16) と EfficientDet-D0 を比較すると以下の通りです。

Model mAP FPS Size [MB]
RefineDet512 (VGG-16) 33.0 24.1 140
EfficientDet-D0 33.8 97 16

上記 FPS は、測定環境が異なるのでご注意ください。RefineDet 512 (VGG-16) では TitanX、EfficientDet-D0 では Tesla V100 による結果です。

mAP では EfficientDet-D0 に軍配が上がりますが、小さい物体に対する AP は RefineDet512 (VGG-16) に軍配が上がります。RefineDet は 2017/11/18 登場と、EfficientDet の三年前なので大健闘とも言えるでしょう。

Model mean S M L FPS
RefineDet512 (VGG-16) 33.0 16.3 36.3 44.3 24.1
EfficientDet-D0 33.8 12.0 38.3 51.2 97.0
EfficientDet-D1 39.6 17.9 44.3 56.0 74.0
EfficientDet-D2 43.0 22.5 47.0 58.4 57.0

EfficientDet の実装

公式 TensorFlow 実装 は 2020/03/14 に最初のコミットがありました。

PyTorch 版は rwightman/efficientdet-pytorch: A PyTorch impl of EfficientDet faithful to the original Google impl w/ ported weights が有力だと思います。mAP がかなりオリジナルに肉薄していますし、混合精度による訓練にも対応しているようです。

PyTorch 版は他にも zylo117/Yet-Another-EfficientDet-Pytorch: The pytorch re-implement of the official efficientdet with SOTA performance in real time and pretrained weights. などがありますが、mAP は上記に比べるとオリジナルとの乖離が大きいです。

EfficientNet

EfficientNet は、2019/05/28 に arXiv に論文が投稿され、2019/07/02 に Google Developers に日本語記事が投稿されました。こちらのすでに多くの解説記事があるので詳細は省きますが、コンピュータビジョンにおけるさまざまなタスクで Backbone の選択に迷ったらとりあえず ResNet を試す、の代替として検討してみてもよいニューラルネットワークモデル、として捉えています。EfficientDet と同様に、速度と精度のトレードオフを調整しやすく、EdgeTPU にも載せられる汎用性の高いモデル、という印象を持っています。

TensorFlow Lite

TensorFlow Lite は、デバイス上での推論 (on-device inference) を可能にするオープンソースの深層学習フレームワークです。Google が開発しており、TensorFlow がベースになっています。 を見ると、Android/iOS/Raspberry Pi などで動作するようで、画像分類物体検出人物姿勢推定音声認識ジェスチャー認識スマートリプライ領域分割画風変換数字分類テキスト分類質問応答などが挙げられています。

TensorFlow Lite のターゲットはモバイルと IoT デバイスで、これらでの低遅延・軽量な推論を目指しています。TensorFlow Lite InterperterTensorFlow Lite Converter という2つの主要コンポーネントで構成されています。エッジで推論を行うことで、以下が改善されるとしています(これらはエッジで推論を行う上で一般的なことで TensorFlow Lite に特化したことではなりません)。

  1. 遅延: サーバーとのデータのラウンドトリップがない
  2. プライバシー: 端末からデータを送信する必要がない
  3. 接続性: インターネット接続を必要としない
  4. 消費電力: ネットワーク接続は電力をかなり消費する

コアは C++ で開発されており、C++ API を基本としていますが、Android では Java、iOS では Swift/Objective-C、もしくは Python からも利用可能になっています。

なお、NVIDIA GPU を活用した高速推論はサポートしていません。NVIDIA GPU を活用したいなら、TensorFlow Lite ではなく TensorFlow を使いましょう。

私が TensorFlow Lite で好きなのは、CPU での推論をシンプルに行えるという点です。バイナリが小さいのも気に入っています。また、XNNPACK を有効にすると、CPU でもこんな速度で推論できるのか、と驚くほど高速化されます。

XNNPACK

XNNPACK は、ARM・WebAssembly・x86 向けの浮動小数点精度でのニューラルネットワーク推論オペレータの最適化ライブラリです。深層学習の研究者などが直接使うようなものではなく、高レベルな機械学習フレームワーク(TensorFlowLite、TensorFlow.js、PyTorch、MediaPipe など)を加速するために低レベルなパフォーマンス・プリミティブを提供します。

2019/09/28 にオープンソースとして公開されたようです。初回コミットがものすごく大きいので、内部ではもっと前から継続的に開発されていたと思われます。また、Facebook の QNNPACK をベースにしているようです3

現状では、以下のアーキテクチャをサポートしているようです4。よくある CPU はおおよそいけるはずです。README には記載されていませんが、ビルドすれば Windows でも動作するようになっています5。Windows での動作には Golang などで有名な mattn 氏が寄与されたようです。

  • ARM64 on Android, Linux, and iOS (including WatchOS and tvOS)
  • ARMv7 (with NEON) on Android, Linux, and iOS (including WatchOS)
  • WebAssembly MVP
  • WebAssembly SIMD (experimental)
  • x86 and x86-64 (up to AVX512) on Android, Linux, macOS, and iOS simulator

TensorFlow Lite からは、XNNPACK はデリゲートとして差し込むことが可能です。

XNNPACK は、TensorFlow 同様 NHWC 形式のテンソルを対象としています。PyTorch などでよく見かける NCHW 形式ではないので注意してください。

CMake によるビルドもできますが、Bazel が本流となっているようで CMake 側は追従が遅れているようです。WebAssembly SIMD も Bazel を参考に CMakeLists.txt を修正すればビルドできます。

WebAssembly SIMD

WebAssembly SIMD はまだ提案段階ですが、Chrome などにはすでに実装されており、「WebAssembly SIMD support」フラグを有効にすると利用可能です。WebAssembly についての説明はここでは省きますが、caniuse.com によると、いつの間にか IE などをのぞく九割のブラウザがサポートしているようです。全体感を俯瞰するため、Fast, parallel applications with WebAssembly SIMD の一部を意訳しておきます。

SIMD は Single Instruction, Multiple Data の略です。SIMD 命令は、複数のデータに対して同じ操作を同時に実行することで、アプリケーションにデータ並列性をもたらす特殊な種類の命令です。音声・動画コーデック、画像処理のような計算が激しいアプリケーションは、SIMD 命令で性能を向上する恩恵を得られるアプリケーションの例となります。最近のアーキテクチャのほとんどは、SIMD 命令を何かしらサポートしています。

WebAssembly SIMD proposal は、最近のアーキテクチャのほとんどで利用可能な SIMD 操作の移植可能で高性能なサブセットを定義しています。この提案は SIMD.js proposal から多くを派生していて、元々は Dard SIMD 仕様から派生しています。SIMD.js proposal は、SIMD 計算を実行するための新しい型と関数をもって TC39 で提案された API でしたが、WebAssemly で SIMD 操作をより透過的にサポートするためお蔵入りとなりました。Web Assemly SIMD proposal は、基礎となるハードウェアを使用してデータレベルの並列処理をブラウザが利用する方法として導入されました。

WebAssembly SIMD proposal の高水準な目標は、WebAssembly Specification にベクトル操作を導入し、移植性のある性能を保証することです。

SIMD 命令集合は巨大で、アーキテクチャによってさまざまな違いがあります。WebAssembly SIMD proposal に含まれている操作集合は、さまざまなプラットフォームで十分にサポートされていて、性能が証明されている操作で構成されています。このため、現在の提案は 128 ビット固定幅の SIMD 操作の標準化に限定されています。

現在の提案は、新しい v128 という値型と、この型に対する新しい操作をいくつか導入しています。操作の選定基準は以下のとおりです。

  • 複数の最近のアーキテクチャでよくサポートされている
  • 命令グループ内の複数の関連するアーキテクチャで性能が向上する
  • 選択した操作集合により、性能断絶を最小限にできる

この提案は活発に開発中で、V8 とそのツールチェインの両方で実験のためのプロトタイプが実装されています。これらはプロトタイプ実装であるため、新しい操作が提案に追加されると変更される可能性があります。

WebAssembly SIMD のユースケース

WebAssembly SIMD proposal は、音声・動画コーデック、画像処理、暗号処理といった負荷の高い計算を必要とするアプリケーションの高速化を目指しています。現在、WebAssembly SIMD は HalideOpenCV.jsXNNPACK のような広く利用されているオープンソースプロジェクトで実験的にサポートされています。

Google Research チームから、いくつか興味深いデモが MediaPipe プロジェクト で公開されています。

SIMD の性能の違いを簡単に観察できる最も視覚的に魅力的なデモのひとつは、手の追跡システムです。SIMD なしでは最近のラップトップで 3 FPS 程度ですが、SIMD ありだと 15 ~ 16 FPS でより滑らかな体験ができます。

SIMD を有効にした Chrome Canary でデモ にアクセスして試してみてください。OpenCV も興味深いです。

WebAssembly SIMD proposal の将来

現在の SIMD proposal は Phase 2 の状態にあり、今後の作業は標準化プロセスで提案を進めることです。固定幅の SIMD はスカラーよりも大幅に性能が向上しますが、最新のハードウェアで利用可能なより幅の広いベクトル操作を効果的に活用していません。現在の提案が進んでいけば、ここで直面する今後の作業は、より幅の広い操作で提案を拡張する可能性を決めることです。

実験

諸事情(コードが汚い、など)により再現コードは掲載できませんが、大まかな手順は以下のとおりです。

  1. TensorFlow のチェックポイントポイントファイルを TensorFlow Frozen Graph に変換する
  2. TensorFlow Frozen Graph を *.tflite 形式に変換する
  3. EfficientDet の後処理のカスタムオペレータを実装する
  4. 全体の推論プログラムを実装する
  5. WASM SIMD 向けにビルドする
  6. ブラウザで動くように HTML + JavaScript で見た目を実装する

EfficientDet の公式実装には、TensorFlow のチェックポイントファイルが公開されています。まず、これをTensorFlow Frozen Graph に変換します。

次に、TensorFlow Lite で利用できる *.tflite 形式に変換します。EfficientDet は RetinaNet 系の物体検出器なので注意が必要です。TensorFlow Lite は TFLite_Detection_PostProcess というカスタムオペレータにより SSD 系の物体検出器の後処理をモデルに入れ込むことができますが、RetinaNet 系ではこのテクニックが使えません。そのため、TFLite_EfficientDet_PostProcess という独自のカスタムオペレータを追加して、それに後処理をやらせることにしました。カスタムオペレータへの入力は、信頼度、位置、それとアンカーとし、出力は TFLite_Detection_PostProcess と同様にしました。独自のカスタムオペレータなので、もちろん推論側でこのカスタムオペレータの処理を実装する必要がありますが、Python 実装はあるので C++ に置き換えるだけで大丈夫です。TensorFlow Lite 形式への変換は、公式実装にある model_inspect.py に手を加えて対応しました。

C++ による TFLite_EfficientDet_PostProcess の実装は、公式実装を見ながら移植していくだけです。以下のように実装しました。

  1. 信頼度に対して Top-K を実行して処理対象を絞り込み、そのインデクスを保持する
  2. 絞り込んだ信頼度をアドレスが連続になるように配置する (Gather)
  3. 絞り込んだインデクスに対応する位置をデコードしつつ、アドレスが連続になるように配置する (Gather)
  4. NMS で検出結果をまとめる

これらの手順を物体のクラスごとに行います。SSD 系と RetinaNet 系では、上記の「位置のデコード」が少し異なるんですよね。

あとは、OpenCV で動画を読み込みながら、フレームを BGR から RGB に変換し、8 bit 符号なしから 32 bit の浮動小数点に変換し、前処理を適用し、推論を実行し、結果を OpenCV で可視化するようにしました。

ここまでで、全体で 1000 行足らずのコードになりました。

あとは流れるように WASM + SIMD で動くように整えていくだけです。Emscripten でビルドできるように CMake の設定を調整しました。TensorFlow Flite は Bazel を元に CMake でビルドできるように調整しました。本筋ではありませんがここが一番時間がかかった気がします。TensorFlow や XNNPACK のディレクトリ構成の把握にもつながったので、無駄ではなかったと信じています。

冒頭にも書きましたが、MacBook Pro (Retina, 15-inch, Mid 2015) の 2.5 GHz で、EfficientDet-D0 を 256x256 の入力解像度の画像を与えたときの速度をベンチマークしました。

ベンチマークの指標として、FPS を計測しました。FPS は画像をモデルに入力してから出力されるまでの時間の平均から算出しています。この時間には、HWC を NHWC を変換し、平均を引いたり、入力画像の解像度をモデルが受け取れる適切なサイズに Bilinear フィルタでリサンプリングするという前処理の時間も含んでいます。入力解像度をのぞき、以下と同じ条件です。

TensorRT で物体検出・姿勢推定はどれくらい速くなる? - OPTiM TECH BLOG

ベンチマーク結果は、以下のようになりました。

Case XNNPACK XNNPACK スレッド数 FPS
A OFF N/A 4
B ON 1 15
C ON 4 20

Chrome で WASM SIMD を有効にしたときの速度は、XNNPACK あり、スレッド数 1 で 7 FPS 程度でした。上記 Case B の約半分です。WASM SIMD では 128 ビット幅しか対応していないことが大きいと思われます。これが 256 ビット幅に対応すると、Case B の 15 FPS にかなり近づいていくと思われます。

まとめ

Chrome で WASM SIMD を有効にすると、ブラウザでもでもいくつかのモデルが現実的な速度で推論できるようになってきています。実装状態を見ると、V8 と SpiderMonkey 以外の対応は進んでいませんが、Edge も Chromium ベースになってきていますし、早く普及が進むことを祈ります。

上記には記載していませんが、顔検出器として有名な RetinaFace は同環境で 35 〜 40 FPS で推論できることを確認しています。

オプティムでは、こうした技術に興味がある・作ってみたい・既に作っている、というエンジニアを募集しています。興味のある方は、こちらをご覧ください。

www.optim.co.jp


  1. https://github.com/google/automl/tree/master/efficientdet のサンプルとして公開されている https://storage.googleapis.com/cloud-tpu-checkpoints/efficientdet/data/video480p.mov を利用させていただきました。

  2. Faster R-CNN、SSD (Single Shot MultiBox Detector)、 YOLO、RetinaNet、Objects as Points などが有名だと思います。

  3. 現在は逆に Facebook が XNNPACK を PyTorch に取り込んでいます。

  4. https://github.com/google/XNNPACK#supported-architectures

  5. https://github.com/google/XNNPACK/commit/462be05c791ee52a0cb279e2cb8ece22783ddd33