メルカリのエッジAIチームのテックリード

大嶋悠司氏(以下、大嶋):『使われる機能を目指して測ったり試したり』という内容で発表していきます。私は運営をしている上田さんと同じメルカリでエッジAIチームというチームのテックリードをしています。

メルカリに入る前は、OSS活動を中心にしていて、例えばさっきも話に出ましたKubeflowの中でハイパーパラメータチューニングをするKatibのオーナーをしていました。メルカリに入ってからは、OSS活動というよりかはサービスよりの開発をしています。

私のチームはエッジAIというチームでして、このチームで行った活動を例に、エッジAI技術を使ったアプリについてお話をしていきたいと思います。

Edge AIとは

まずエッジAIとは? という話ですね。このエッジAIというワードは最近ちらほらよく聞くようになってきたので、ご存知の方もけっこういるかもしれません。

ふつう機械学習のサービスを立てるときは、なにかしらクラウド上なりにリッチなサーバーを置いて、そこにWeb REST APIなりgRPCのAPIなりを開けて、端末からアクセスするとかブラウザからアクセスするっていうようなことが多いと思います。

それとは違って、例えばスマホ上で機械学習のモデルを搭載して、インターネットを介してアクセスすることなく端末上で機械学習の推論を行うような技術。まあ推論だけではなく学習を行うこともありますが、そういった技術になります。

Edge AIでできること

じゃあそれを使って何が嬉しいのか。何ができるのかという話ですね。想像できると思いますが、インターネットを介さないので非常にレイテンシーが低いですね。これ以外にもメリットはあるんですが、我々は主にここを重視しています。

例えばインターネットを介してサーバーにリクエストを送る場合だと、モデルをどれだけチューニングしたとしてもそもそもインターネットを介さないといけないので、まあ500ミリsecを切ることは難しいんじゃないかと私は想像します。1秒とか2秒かかるようなサービスもぜんぜん普通にあると思います。

それに対してEdge、つまり端末上で機械学習のモデルを動かしたときには、モデルをチューニングさえすれば50ミリsec以内、早ければ10ミリsecとかで返すことも可能ですね。ただデバイス上で動かさないといけないので、モデルサイズだったり消費電力に制限がある技術となります。

じゃあ一体我々は何のためにそれに取り組んでいるのか? という話ですが、いろいろ制限は大きいんですが、まずレイテンシーが非常に低いこと。インターネットを必要としないで推論できる、つまりオフラインでも推論ができるということで、インターネット環境に影響されずに安定して推論が行えることから、それを使ってUXを大きく改善していけるだろうということで、この領域の技術に私たちは注力しています。

Edgeっていうと実はいろいろあるんですね。製造業だと、例えば工場のラインにIoTデバイスを置いてその上で推論をさせるとか、そういうエッジAI用のチップとかも売られてたりします。メルカリは今のところ、そういったデバイスを作ってなくて、スマホ上で機械学習を行うということを指してエッジAIというふうに呼んでいます。

エッジAIで使うフレームワークと構成

機械学習をモバイル上で行うというのはけっこう最近流行っている技術でもあって、それを使うためのフレームワークがいくつかあります。TensorFlow Lite、TVM、PyTorch Mobile。こういったものを使って機械学習を端末上に置いていきます。

ただやることはどのフレームワークも基本的には同じで、まず普通に機械学習のモデルを学習します。そのモデルを端末上に置けるように軽量化していきます。例えば量子化をしたりだとか、プルーニングをしたり。このへんはあとで詳しく話します。

モデルの形式を変換したあとデバイス上で実行する。これらを3つのフレームワークのどれかを使ってやります。我々は2020年5月現在、最も成熟しているTensorFlow Liteを使って、このエッジAI技術をアプリに組み込むことにしました。

エッジAIで作ったもの

では我々は何を作ったかという話です。メルカリのアプリを使ったことがある方だとご存知かもしれませんが、メルカリって、商品を出品するときに、例えば本とかDVDとかそういったものだとバーコードを読み込むとそのバーコードの情報がサッと入力されます。

商品のタイトルとか詳細とかに勝手に入れてくれる便利な機能があります。けっこう使っていただいてはいるんですが、もっと認知を広めていきたいという思いがあります。

じゃあこのバーコード出品をもう少し認知率を高めるためにはどうしたらいいか?  ということで、バーコード出品を使わずに出品しようとした場合、例えば本なんかを普通に写真で撮影したときに「これは本だからバーコード出品したらいいんじゃないですか?」みたいなサジェストを行えれば、もうちょっとバーコード出品を知らない人に知ってもらえるんじゃないかと考えました。

「本じゃないですか?」って判断するだけなら、実装は、MobileNetみたいなImage Classificationモデルで普通にできそうですよね。ということで、じゃあ本だと認識したらサジェストを出すという機能を作ってみようということになりました。この機能については、先日Googleのテックブログに記事を書かせていただきましたので、そちらと今回の発表は少し重複する部分もあります。

TensorFlow Liteを使ったモデル学習

ではこの機能を作っていく流れで、先ほどのTensorFlow Liteを使ってどうやって開発していったかを説明していきたいと思います。

まずモデル学習です。モデル学習は基本的には普通にMobileNet V3を学習しました。ただその中でQuantization Aware Trainingというものを使います。これはこのあとでちょっと説明するんですけど、モデルを軽量化するために量子化っていうのを行うんですね。この量子化をしたときに精度劣化を防ぐための技術です。

普通ニューラルネットワークの重みはFLOAT32で記述されているんですけれども、それをint8の8bitに落とし込む。量子化します。そうすることで単純に情報量を減らし、モデルのサイズを小さくすることが可能になります。

ただもちろんそうすると情報量が減っちゃうので精度が劣化します。それを防ぐために、こいつは将来quantizeされる、量子化されるんだぞということを意識して学習の重みを考えて配置していくというような学習方法です。

これを使わずにquantizeすると精度が5ポイントくらい落ちたのに、使うことでほぼ落ちないという効果があったので、これだけちょっと学習のときは工夫をしたという話になります。

モデルの軽量化と形式の変換

次にモデルの軽量化とか、そのあとモデルの形式の変換です。モバイル端末上で推論を行うために必要なことは何かというとですね。

今回のアプリでは、写真を撮った瞬間にサジェストを出したいわけですよね。必要とされる推論速度についていうと、写真を撮って1秒後とか2秒後とかだったら、もうカメラを閉じて次の出品行動をとって、タイトルを入力するとかに移っている可能性が高いのでそれだと困る。即座に出したい。だから50ミリsec以内に推論を返したいと。

さらに機械学習のモデルを端末上に入れないといけない。そうなると機械学習のモデルを入れるからアプリのサイズが100メガとか300メガ大きくなりますっていうのは基本的には許されないので。10メガよりも小さくしてくださいねっていうことでやっています。

TensorFlow Liteでモデルを変換するときにオプションが4つあります。これで先ほどから何度も出てきている量子化というのを行います。

選択肢として、FLOAT32が量子化しない。普通にFLOAT32で重みを表現する。FLOAT16で重みを表現する。例えばGPUで、FLOAT16で計算できるものがありますけれども、あれと同じようにFLOAT16で表現する。そうするとモデルサイズは単純計算で半分になりますね。

8bitで表現する。その中でも重みだけを8bitで表現するものと、重みと活性化関数両方を8bitで表現できるようにするものの2種類があります。

これはTensorFlow Liteの仕様によるところが大きいんですけれども、実は重みだけを量子化するのが、ファイルサイズが一番小さいです。活性化関数まで量子化しようとすると活性化関数を8bitで表現するために、この値が入ってきたらこれを返すみたいなマップを全部保持することになるので、モデルサイズが少し大きくなるというような仕様になっています。

レイテンシーについては、活性化関数まですべて8bitで表現できていれば計算時も8bitで計算するので非常に速い。しかし重みだけ8bitで表現されている場合は計算時にもう1回FLOATに直して計算するので別に速くはならないという結果になっています。

ただ今回は重みだけの量子化でレイテンシーを十分達成できるため、重みだけの量子化を選択しました。活性化関数まで量子化をしようとすると、実はいろいろ大変なことがあるので、またそれは別のところで話せればと思います。

ランタイムとデバイス上での実行

最後にデバイス上で実行するにあたってです。ここでもTensorFlow Liteのオプションがいろいろあります。TensorFlow Liteで普通にCPU上で実行するというのが一番普通のやつなんですが……。

それを例えばGPUで計算を行わせるようなGPU delegationだったりとか。iOSだったらCoreMLのライブラリを利用するCoreML delegationとか、ほかにもいくつかオプションがあるんですね。

これらを行うとパフォーマンスが良くなったり良くならなかったりというようなものになるんですが。今回はdelegationなしでもパフォーマンスが十分出るということがわかったので、delegationはまた今後やっていこうねっていうことになっています。なので基本的なTensorFlow Liteの機能だけを使ってモデルを作りました。

アプリの完成。そしてこれから

これができた。完成! 写真をカメラで撮ると「これ本じゃないですか?」っていうサジェストが出るっていう機能になります。お疲れ様でした! できました! これで完成ですね。アプリをリリースしますね。終わりですね。

たぶんMLCTに来られているみなさんは「いや、嘘でしょ」っていうのがわかると思います。

タイトルにも書きましたとおり、ここからが本題というか、ここからも大事な場面ですよね。