姿勢推定モデル MoveNet を TensorRT でベンチマーク

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 のモデルを入手できます。

• movenet/singlepose/lightning v4
• movenet/singlepose/thunder v4

TFLite モデルを ONNX モデルに変換

まず、ベンチマークのために TFLite モデルを ONNX モデルに変換します。ONNX モデルしてしまえば、ONNX Runtime で CPU/GPU でベンチマークできますし、TensorRT でもベンチマークできるので便利なのです。最近の世の中は便利なもので tf2onnx を使うと TFLite モデルが簡単に ONNX モデルに変換できます。例えば、物体検出器の MobileDet なら以下のとおりです。

• Convert a mobiledet tflite model to ONNX

上記のように、環境さえ作ってしまえば簡単に変換できます。今回は以下のように環境を用意しました。

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 実装などが登場するとありがたいですね。

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

www.optim.co.jp