ニューラルネットワークを疎にして推論を高速化(1/2)

１．ニューラルネットワークを疎にして推論を高速化(1/2)まとめ

・モデル最適化には大部分の重みが0に設定されているスパースニューラルネットワークが有効
・しかし良く使われる畳み込みスパース化するツールとサポートするツールがなかった
・XNNPACKアクセラレーションライブラリとTensorFlow Liteのアップデートで可能になった

２．スパースニューラルネットワークとは？

以下、ai.googleblog.comより「Accelerating Neural Networks on Mobile and Web with Sparse Inference」の意訳です。元記事の投稿は2021年3月9日、Artsiom AblavatskiさんとMarat Dukhanさんによる投稿です。

「密です！」は少し前に流行りましたが「疎ら(まばら）です！」は漢字表記はあまり見慣れないですね。dense(密)モデルとsparse(疎)モデルはセットで覚えてしまうと良いと思います。

アイキャッチ画像のクレジットはPhoto by Robert Anasch on Unsplash

ニューラルネットワークをスマートフォンなどのデバイス上で実行する事により、低遅延でプライバシーを意識した方法で、人間のポーズ推定や背景のぼかしなどのさまざまなリアルタイムアプリケーションが可能になります。

エンジニアは、XNNPACKMLアクセラレーションライブラリを備えたTensorFlow LiteなどのML推論フレームワークを使用して、モデルサイズ、推論速度、予測品質、の３点のバランスが取れるスイートスポットを見つけることにより、様々なデバイスで実行可能なようにモデルを最適化します。

モデルを最適化する1つの方法は、重みのかなりの部分がゼロに設定されているスパースニューラルネットワーク(訳注：sparse neural networks、レイヤー間の接続を間引いてまばらにしたネットワーク。行列内の0が増えると言う事なので行列演算がとても簡単になります)を使用することです。

一般に、これは、圧縮によってモデルサイズを縮小するだけでなく、積和演算操作のかなりの部分をスキップする事で推論を高速化できるため、望ましい品質をもたらします。

更に、モデル内のパラメーターの数を増やしてから、元のモデルの品質に一致するようにモデルをスパース化すると、アクセラレータを使った推論速度向上の恩恵も受けることができます。

ただし、この手法の使用は、良く使われる畳み込みアーキテクチャをスパース化するツールがないこと、およびこれらの操作をデバイス上で実行するためのサポートが不十分であるため、本番環境では制限されたままです。

本日、XNNPACKアクセラレーションライブラリとTensorFlow Liteの一連の新機能のリリースを発表しました。

これにより、スパースネットワークの効率的な推論が可能になり、ニューラルネットワークをスパース化する方法のガイドラインとともに、研究者が独自のスパースなオンデバイスモデルを開発できるようになります。

DeepMindと共同で開発されたこれらのツールは、MediaPipeのハンドトラッキングやGoogle Meetのバックグラウンド機能など、新世代のライブ知覚体験を強化し、モデルサイズを半分に減らしながら、推論速度を1.2倍から2.4倍に加速します。

本投稿では、スパースニューラルネットワークの技術的な概要(トレーニング中のスパース性の誘発からデバイスへの展開まで)を提供し、研究者が独自のスパースモデルを作成する方法についていくつかのアイデアを提供します。

Google Meetのバックグラウンド機能の同じ品質のdense(密)モデルとsparse(疎)モデル(右)の処理時間の比較。わかりやすくするために表示されている上部の処理時間は100フレームにわたる移動平均です。

ニューラルネットワークのスパース化
MobileNetやEfficientNetLiteなどの多くの最新の深層学習アーキテクチャは、主に、小さな空間カーネルを使用した深さ方向の畳み込みと、入力画像の特徴表現を線形に組み合わせる1×1の畳み込みで構成されています。

このようなアーキテクチャには、多くのネットワークの開始時に頻繁に発生する完全な2D畳み込みや深さ方向の畳み込みなど、スパース化の潜在的なターゲットがいくつかありますが、推論時間で測定した場合、最もコストのかかる演算子は1×1畳み込みです。コンピューティング全体の65%以上を占めるため、スパース化の最適なターゲットです。

Architecture	Inference Time
MobileNet	85%
MobileNetV2	71%
MobileNetV3	71%
EfficientNet-Lite	66%

最新のモバイルアーキテクチャの1×1畳み込み専用の推論時間の%単位の比較

XNNPACKなどの最新のオンデバイス推論エンジンでは、深層学習モデルでの1×1畳み込みの実装やその他の操作は、HWC(Height, Width, Channel size)テンソルレイアウトに依存しています。HWCテンソルレイアウトでは、テンソルの寸法は、入力画像の高さ、幅、およびチャネル(例えば、赤、緑、または青)に対応します。

このテンソル構成により、推論エンジンは、各空間位置(つまり、画像の各画素)に対応するチャネルを並列に処理できます。ただし、このテンソルの順序は、チャネルをテンソルの最も内側の次元として設定しており、アクセスするための計算コストが高くなるため、スパース推論には適していません。

XNNPACKのアップデートにより、モデルがスパースであるかどうかを検出できるようになります。その場合、標準の密推論モードから疎推論モードに切り替わり、CHW(チャネル、高さ、幅)テンソルレイアウトが採用されます。

テンソルのこの並べ替えにより、2つの理由でスパース1×1畳み込みカーネルの実装を高速化できます。

(1)対応するチャネルの重みがゼロの場合、画素毎にチェックせずとも、ゼロか否かをチェックするだけでにスキップできるのでテンソルの空間操作が簡単になります。

(2)チャネルの重みがゼロ以外の場合、隣接する画素を同じメモリユニットにロードすることで、計算をより効率的に行うことができます。

これにより、複数の画素を同時に処理しながら、複数のスレッド間で各操作を並行して実行できます。これらの変更を合わせると、重みの少なくとも80%がゼロの場合、1.8倍から2.3倍の速度向上になります。

スパース推論に最適なCHWテンソルレイアウトと各操作後に標準的なHWCテンソルレイアウトに変換する事を回避するために、XNNPACKはCHWレイアウトで複数のCNN演算の効率的な実装を提供します。