R&D チームの奥村（@izariuo440）です。日本で気温が暖かくなると姿勢推定が盛り上がってくるように見えるのは私だけでしょうか？今年の５月に発表された人物姿勢推定モデルの MoveNet v3 をTensorRT でベンチマークしてみました。CPU でもかなり高速推論が可能ですが、CUDA でもっと速く、TensorRT でもっともっと速くなることが確認できました。
MoveNet
MoveNet は高速・軽量・高精度と三拍子揃った人物姿勢推定器です。Lightning/Thunder の二種類があり、前者はより軽量・高速です。詳細は MoveNet.SinglePose Model Card (PDF) や Next-Generation Pose Detection with MoveNet and TensorFlow.js — The TensorFlow Blog が参考になると思います。
TensorFlow Hub から TensorFlow SavedModel、TensorFlow.js、TFLite のモデルを入手できます。
TFLite モデルを ONNX モデルに変換
まず、ベンチマークのために TFLite モデルを ONNX モデルに変換します。ONNX モデルしてしまえば、ONNX Runtime で CPU/GPU でベンチマークできますし、TensorRT でもベンチマークできるので便利なのです。最近の世の中は便利なもので tf2onnx を使うと TFLite モデルが簡単に ONNX モデルに変換できます。例えば、物体検出器の MobileDet なら以下のとおりです。
上記のように、環境さえ作ってしまえば簡単に変換できます。今回は以下のように環境を用意しました。
pip install tensorflow==2.5.0 pip install onnxruntime==1.8.0 pip install onnx==1.9.0 pip install tf2onnx==1.9.0
次に、ダウンロードしたモデルを変換します。今回は最新の v4 ではなく v3 を使いました。--opset は用途に応じて調整してください。
python -m tf2onnx.convert --opset 11 \
--tflite movenet-lightning-v3.tflite \
--output movenet-lightning-v3.onnx
ONNX モデルを加工
先ほど得られたモデルには、合計 212 のノードが含まれており、オペレータは 21 種類あります。
| 名前 | 出現階数 | 出現割合 [%] |
|---|---|---|
| Add | 19 | 9.0 |
| ArgMax | 2 | 0.9 |
| Cast | 8 | 3.8 |
| Clip | 35 | 16.5 |
| Concat | 7 | 3.3 |
| Conv | 74 | 34.9 |
| Div | 3 | 1.4 |
| GatherND | 3 | 1.4 |
| Mul | 7 | 3.3 |
| Relu | 7 | 3.3 |
| Reshape | 13 | 6.1 |
| Resize | 3 | 1.4 |
| Shape | 3 | 1.4 |
| Sigmoid | 2 | 0.9 |
| Slice | 3 | 1.4 |
| Split | 2 | 0.9 |
| Sqrt | 1 | 0.5 |
| Squeeze | 4 | 1.9 |
| Sub | 6 | 2.8 |
| Transpose | 4 | 1.9 |
| Unsqueeze | 6 | 2.8 |
この中に GatherND があるのですが、TensorRT はこのオペレータには対応していないため、このままでは TensorRT でベンチマークできません。MoveNet のケースでは GatherND を Gather/Reshape/Shape/Div/Slice の組み合わせで置き換えることができました。他にもいくつかの辻褄合わせをすることで TensorRT の nvonnxparser が解釈可能な ONNX モデルへと加工することができます。具体的な手順はそれなりに複雑なのでここでは割愛します。
ベンチマーク結果
先ほど得られたモデルを FP32 の精度でベンチマークすると、以下の結果が得られました。
| モデル | 入力解像度 | 容量 | A | B | C | D |
|---|---|---|---|---|---|---|
| MoveNet Lightning | 192x192 | 約 10 MB | 200 FPS | 210 FPS | 340 FPS | 740 FPS |
| MoveNet Thunder | 256x265 | 約 25 MB | 60 FPS | 80 FPS | 220 FPS | 460 FPS |
A/B/C/D はそれぞれベンチマーク環境で、以下のとおりです。
- A: MacBook Pro (Retina, 15-inch, Mid 2015)、CPU 2.5 GHz クアッドコアIntel Core i7、ORT v1.8.0、CPU
- B: Intel Core i7-7700 @ 8x 4.2GHz、GeForce GTX 1080 Ti、ORT v1.6.0、CPU
- C: 同上、ORT v1.6.0、CUDA
- D: 同上、TensorRT Release 21.05
TensorRT がかなり高速ですね。FP16 ではさらに速くなるはずです。
まとめ
MoveNet の TensorRT ベンチマークは見たことがなかったので確認してみました。MoveNet も PyTorch 実装などが登場するとありがたいですね。
オプティムでは、こうした技術に興味がある・作ってみたい・既に作っている、というエンジニアを募集しています。興味のある方は、こちらをご覧ください。
