Intel GPUでJPEGデコード・推論

動機など

大量の画像ファイルに対して画像認識とかをしたいときがありますよね。

• 画像のデータセットに前処理をしたい
• 自分の撮った写真を分類したい
  などなど....

普通に実装しようと思ったら

という風にやるのが典型だと思います。
実際にやってみると、JPEGのデコードがボトルネックになってGPUを使いきれないということがわかってきます。
JPEGのデコードって重い処理なんですよね...
NVIDIA GPUならtorchvision.io.decode_jpegという関数があって、これでNVJPEGを呼べます。
これによって、CPUはファイルを読んでGPUに投げ、(場合によっては)後処理をするだけで画像認識ができます。

で、タイトル通り、今回はこれをIntel GPUでもやりたいというわけです。

QSVの実装をする

Intel® Video Processing Library というものがあって、ffmpegのQSVはこれを叩いているので、これをC++から使って、nanobindでPythonのライブラリを自作すると良さそうです。
全体像はソースコードを読んでいただきたいのですが、ざっと

MFXCreateSession //ライブラリの初期化
MFXVideoDECODE_DecodeHeader //JPEGのヘッダ解析
MFXVideoDECODE_Init //JPEGのヘッダの情報に基づいて、Media Engineを初期化
MFXVideoDECODE_DecodeFrameAsync //ビットストリームを投げてエンコード開始

という感じで処理を行います。
一つ問題があって、JPEGを毎回デコードするごとにMedia Engineの初期化が必要です。
半年前くらいまでめちゃくちゃ遅かったのですが、ドライバを最新に更新したら、ある程度は使える速度になっていました。

これがJPEGの処理1回分をズームしたVTuneのキャプチャなのですが、ID3D11VideoDevice::CreateVideoDecoderが時間を食っていることが分かります。一般に動画をデコードする場合、動画の最初のフレームを処理する際に1度だけMedia EngineをID3D11VideoDevice::CreateVideoDecoderで初期化してやればよいですが、今回はJPEGファイルを処理するため、毎ファイルごとにハフマンテーブルや解像度が異なっているので、初期化してやる必要があります。
実際、ID3D11VideoDevice::CreateVideoDecoderのドキュメントを読むと、引数のD3D11_VIDEO_DECODER_DESCのメンバに画像の高さと幅を要求しています。VPLライブラリがバックエンドでDirectX11を使うのをやめないとこれ以上速くならなさそうですが......

nanobindでPythonに渡す

C++側

nanobindは、Pythonのライブラリを作るのを簡単にしてくれるライブラリで、よく知られたpybind11よりもシンプルに使えるところがウリらしいです。
nanobindはNumPyとかDLPack対応のデータ形式を渡せるらしいですが、自分がやってみた限りだとbad castという実行時エラーが出てどうしようもなかったので、ポインタ渡しをすることにしました。nanobindにポインタ渡しという概念は無いので、とりあえず数値に変換して渡してやることにします。C++側ではreinterpret_cast<intptr_t>(ptr)するだけです。

Python側

Python側では、ctypesでガチャガチャやることによってC/C++のポインタを読むことができます。

y_arr = numpy.frombuffer((ctypes.c_uint8 * (pitch_h * pitch)).from_address(ptr), dtype=numpy.uint8,
                         count=pitch_h * pitch)
print(f"{y_arr=}")
uv_arr = numpy.frombuffer((ctypes.c_uint8 * (int(pitch_h * 1.5) * pitch)).from_address(ptr),
                          dtype=numpy.uint8,
                          count=int(pitch_h / 2) * pitch, offset=pitch_h * pitch)
print(f"{uv_arr=}")

(ctypes.c_uint8 * size).from_address(ptr)でポインタが読めるようです。それをnumpy.frombufferで読み込みます。この演算子オーバーロードはさすがにヤバいんじゃないかと思いますが、サンプルにそう書いてあったのでしょうがないです。DLLは同一のメモリ空間なんだなあって実感しますね。

NV12といって、Y(輝度情報)を本来のピクセル分、U(青-輝度)を4分の1、V(赤-輝度)を4分の1ずつ持たせて情報量を削った形式で返されるので、このままでは画像認識モデルに突っ込めないです。そこでRGBに変換する必要があります。

なお、図では、Y1からY48まで連続して得られているように見えますが、実際は一列ごとに右端がpaddingされています。そもそもJPEGが8x8ブロックでMCUを構成しているから、大体8の倍数になっています。
よって、そのままメモリ配列をNumPy配列にしてpyplot.imshowすると、せん断されたような表示になってしまいます。つまり、NumPyで言うところの行列のサイズとストライドが異なっている状態ですね。
NumPyには、numpy.lib._stride_tricks_impl.as_stridedという、ストライドを手動で指定できるありがたい関数があるので、

y_plane = as_strided(y_arr, (pitch_h, pitch_w), (pitch, 1))
uv_plane = as_strided(uv_arr, (int(pitch_h / 2), int(pitch_w / 2), 2), (pitch, 2, 1))

といった感じで読み込めます。uv_planeは、それぞれ縦横2倍にrepeat(2, axis=0).repeat(2, axis=1)することで、y_planeとサイズを合わせます。

Y	U	V	元画像

\begin{bmatrix}
R \\
G \\
B \\
\end{bmatrix}
=
\begin{bmatrix}
1 & 0 & 1.402 \\
1 & -0.344136 & -0.714136 \\
1 & 1.772 & 0 \\
\end{bmatrix}
\begin{bmatrix}
Y \\
U - 128 \\
V - 128 \\
\end{bmatrix}

という行列の積でRGBに変換できます。
自分の扱ってるJPEG画像をざっと調べたところ、大体Rec.601という色空間だったので、上記の式の変換行列を用いればいいことが分かりました。

ycbcr_mat = yuv_plane.transpose((1, 2, 0)) - [0, 128, 128]
# print(ycbcr_mat)
transform_matrix = numpy.array([
    [1, 0, 1.402],
    [1, -0.344136, -0.714136],
    [1, 1.772, 0]
])
rgb_plane = (numpy.clip(numpy.dot(ycbcr_mat, transform_matrix.T), 0, 255).astype(numpy.uint8))

これでOK.
ここまで出来たらあとはOpenVINOで推論しておしまいです。

実行結果

平常時の負荷の目安

1	2	3
Intel QSV
Pillow

結論としては、推論速度を保ちつつCPUの負荷はかなり抑えられる！といった感じでしょうか。
※突っ込んでくる人がいると思うので書いておきますが、Pillow-SIMDはPillowと同等の速度でした。

結論

Intel GPUしか搭載してない環境(Intelのラップトップとか)においては、ハードウェアにオフロードして低コストでそれなりの推論ができると思うが、デスクトップ環境では普通にNVIDIA GPU使うのが速くて便利だと思います。

真似したい変人さんはこことここにソースあるんで読んでみるなり、僕に聞くなりしてください。
あとは、最近Intel® Video Processing Libraryに、デコードした画像をOpenCLのメモリとしてエクスポートする機能が生えたので、この機能を使えば一度もCPU側にデータを返さずにOpenVINOで推論できるんじゃないか？と画策しています...(時間があったらやるかも)