[iOS] エッジMLの前処理・後処理でMetalを利用してみた
概要
iOSアプリ開発で「エッジML」（後述）を行なっているのですが、その前処理・後処理にMetalを利用してみました。
自身が向き合っていた課題に対してMetalを利用するのがベストソリューションだったのかはいまも確信を持てていないのですが、目的としていたことが一定実現できたのと、自分用の備忘録としてもどこかに残しておきたかったのでこちらに書く次第です。
（初めてのzennでの投稿になります🙏）
参考にしたもの
Metalを実装するのが初めてだったので、以下の書籍を購入して勉強しました。
謝辞
はじめにお断りしておくと、上記書籍のサンプルコードをそのまま利用している箇所がそれなりにあることと会社の都合もあわせて、具体的なコードは省略させていただきます。
「どんな課題に対してMetalを利用したのか」「その課題は解決できたのか」「何かハマりどころなどはあったか」「他の技術的な選択肢はあったのか」といった内容を書いていきます。
解決したかった課題
現在私は「ForceSense」という野球のフォーム解析を行うiOSアプリを開発しています。
どんなアプリなのかをもう少し簡単に説明すると、「撮影した動画から選手のフォームを解析・分析し、パフォーマンス向上に役立てる」といったものになります。
フォーム解析[1]については、社内のMLエンジニアに用意していただいたTensorFlow Liteのモデルを用いてアプリ側で推論処理を実行（エッジML）しています。
またその推論処理を行うにあたって必要な前後の処理、いわゆる前処理・後処理もアプリ側で行なっています[2]。
前処理・後処理の具体的なイメージはTensorFlow LiteのPoseNetのexampleが参考になるかなと思います。
前処理
後処理
私自身としては、これらのコードを参考にしつつ今回利用したモデルに合わせて前処理・後処理を実装していくとういことをしました。
さて、本題である「解決したかった課題」ですが、この前処理・後処理ではそれぞれループ処理が数十万回ほど走っていたということもあり、実行時間がそれぞれ0.02−0.035secほどがかかっていました[3]。
0.02secというとそれほど大きな時間ではないのではと思われる方もいるかもしれませんが、先述したように「撮影した動画からフォームを分析」する必要があるため動画の各フレームに対して実行する必要があります。
例えば「3sec/60fps」という動画があるとすると、この動画には180フレーム存在することになるので、「180 * 0.02 * 2 = 7.2sec」もの時間が前処理・後処理だけでかかってしまうことになります。（ここにモデルの推論処理の時間などもかかってきます）。
この前処理・後処理の実行時間をどうにかして短くしたいという課題に対してMetalを使って解決を試みた、というのがこの記事の内容になります。
GPGPUとしてのMetal
Metalというと「描画処理を高速可能にするもの」というイメージが自分としては強かったのですが、GPGPU（General-Purpose computing on Graphics Processing Units）つまり「GPUの演算資源を画像処理以外に応用する技術」[4]としても利用できます。
もちろんGPUによる計算なので得意・不得意があり、「同じことを繰り返し行い、各計算は他の計算結果を利用せず、並列して、順番も関係なく実行できるような計算処理」を得意としています。[5]
今回前処理・後処理が行っていることを簡単に書くと、
- 前処理: 画像のそれぞれの画素のRGBの値を、モデルの入力に必要な値に正規化する
- 後処理: 体の部位ごとに、出力された2次元ベクトルの中から必要な値（インデックス）を求める
といったものになります。これは上記GPUの得意なことに当てはまると思います。
Metalを利用したときのおおまかなコード
先にリンクを貼ったTensorFlow LiteのPoseNetのexampleのようなコード（CPUで実行される）から、Metal（GPUで実行される）に置き換えていくということを今回行ったわけですが、その置き換えられたコードは、あくまで私自身の整理ですが、以下の3つのパートに分かれると思っています。
- Metalを利用するまでのいくつかの処理
- Metalの処理の呼び出し
- Metal（GPU）で行う処理
Metalを利用するまでのいくつかの処理
Metalを利用するまでに、以下のような多くの処理が必要になります。
- GPUを検索する
- コマンドキューを作る
- コマンドバッファを作る
- コマンドエンコーダーを作るパイプラインを作成する
コード行数でいうと30行前後でしたが、私自身はここで何をしているのか深く理解できておらず、この部分は前述した書籍のサンプルコードをほとんどそのまま持ってきた格好になっています。
Metalの処理の呼び出し
Metalの処理を呼び出します。
Metalで処理を行う際にスレッドをどのような形（MTLSize）で行うかを指定するところが慣れるまで時間を要しましたが、そこ以外は特に難しいことはなかったと記憶しています。
Metal（GPU）で行う処理
Metal（GPU）で行う処理は、.metalファイルにて、MSL(Metal Shader Language)でプログラミングします。
MSLはC++14をベースとした言語のようです。
C++はこれまで私自身は書いたことがなかったのですが、そもそもここで書くコードは複雑なものではないはずであること（GPUでは複雑なロジックは扱えない）、またXcode上でプログラミングする際には「補完が効く・コンパイラがエラーを表示してくれ」たので特段困ったことはなかったです。
以下MSLのドキュメントになります。
Metalで課題は解決した？
「で？？結局Metalを使って速くなったの？速くなったのならどれくらい速くなったの？」...一番重要なところです。
結論からいうと、つまりは10倍ちょい速くなった計算です👍👍👍
| before | after | |
|---|---|---|
| 実行時間 | 0.02−0.035sec | 0.001-0.002sec |
特にどれくらい速くなることを目指していたか・期待していたか、というのは決めてなかったのですが、これは今回実現したかったことを十分賄えているなという印象です。
わかっていないこと
まだまだMetalについてはわかっていないことが多々あるのですが、１つだけ、わかっていないことというか「どのように気をつければいいの？」というものを書いておこうと思います。
今回利用したAPIの１つに以下のものがあります（GPUで処理する際のスレッド数を指定するもの）。
上記リンクを開いてもらえればわかりますが、赤く"Warning"と表示されている箇所があると思います。内容をみてみると、
Use this method only if the device your app is running on supports nonuniform threadgroup sizes. Check for device capabilities with supportsFamily(_:) on the device providing your compute command encoder. See Metal Feature Set Tables (PDF) for device support information.
とあります。要は「端末のGPUによってはサポートされていないメソッドなので、supportsFamilyメソッドでチェックしてね」と書かれている。
対応自体はその通りにすれば問題ないはずですが、「このメソッドは特定のGPU Familyでしか動かない」ということに気づく術が、ドキュメントを逐一参照する or 動作検証（しかもできるだけ古い端末を使う）で気づくしかないのが辛いなと思いました（何か良い方法があるのだろうか？）。
サポートされているメソッド（API）とそのGPUのバージョンの対応表は以下にあります。
他のソリューション？
概要にて「今回の課題に対してMetalがベストソリューションだったかわからない」という旨を書きました。
では他のソリューションは？といわれると、少なくともSIMDを利用しても今回の課題については解決できたかもと思っています。
SIMDについては今年（2024年）のiOSDCでも素晴らしいトークがあったのでそちらを観ていただけるといいと思いますが、ざっくばらんにいうと「CPU上で並列計算を行い処理を高速化する」というもの（SIMDはSingle Instruction, Multiple Dataの略）
SIMDでも今回の課題を解決できたかもしれないというのにはそれなりの根拠があって、実は今回Metalに置き換えた箇所以外でも似たような処理を行っていた箇所が１つあった（画像処理において、ループ文で数十万回の処理を行い0.02秒ほどの実行時間がかかっていた）のですが、そちらはOpenCVのメソッドがそのまま利用できることを失念していたのでそちらに置き換えたところ、ほぼ同じ精度で改善が実現できました。
OpenCVは上記トークにもあるのですがSIMDを利用しています。これが「SIMDでも今回の課題を解決できたかもしれない」と思った次第です。
ではMetalとSIMDどちらで対応するのが良かったのかと問われると、Metalを学んだ分のサンクコストを差し引いてもMetalという選択は悪くはなかったと思っています。
理由としては、
- Metalはそれを利用するまでに多くのコードを書かなくてはいけないが、ロジック部分についてはそこまで複雑ではないという印象を持っているため
- （上記トークにて述べられているのですが）SwiftのSIMDはどうやら内部的に未実装のものが多々あるらしく（！？）、今回SIMDで対応しようとしても結局はC++で書かないといけないであろう点
- これは完全に自信がないのですが、愚直なSIMDのコードを書いただけでは処理速度は4倍程度にしかならないのではと考えられること[6]。つまりOpenCV並みの処理速度を実現しようと思うと、もっと複雑なアルゴリズムで最適化しなくてはいけないのではと思われること。ここについてはシンプルにOpenCVのコードを読むなりすればいいのだとは思うのですがやれてないです🙇♂️
-
機械学習における「姿勢推定」「姿勢検出」と呼ばれるものとイコールです。（画像に人が写っているときに、その体の各部位が画像内のどの位置にあるのかを予測するもの） ↩︎
-
Appleの
Visionの姿勢推定を行うVNDetectHumanBodyPoseRequestでは前処理・後処理はモデル内に含まれている（はず）。TensorFlowのモデルでも同じようなことはできると思うのですが、いくつかの事情からアプリ側で前処理・後処理を行っています ↩︎ -
この記事内における実行時間は、あくまでも私個人のiPhone 15 Proの実機で計測したものです ↩︎
-
上記書籍から引用 ↩︎
-
iPhone（ARMプロセッサ）においてSIMDで一度に処理できる最大bit長は128、今回扱っている値がFloat型（32bit）のため（合ってる？？） ↩︎
Discussion