CUDAによるバンドル調整の高速化とソースコード公開

はじめに

こんにちは、エンジニアの高木です。

私は現在、adaskitという社内の自動運転関連のオープンソースプロジェクトに携わっており、プロジェクトのこれまでの成果としてlibSGMやsegmentation-sgm、「Visual OdometryをCUDAで高速化した話」等を紹介しました。

今回はVisual SLAM等で用いられるバンドル調整について、CUDAによる高速化に取り組んだ話を紹介します。結果だけ先に言うと、CPUで約12秒かかる規模のバンドル調整に対し、GPUで約1.2秒まで短縮することができました。

また、そのソースコードをGithubに公開しましたので、興味があれば是非覗いてみてください。
fixstars/cuda-bundle-adjustment

背景

バンドル調整(Bundle Adjustment)とは、画像からの3次元復元において3次元点群とカメラ姿勢を改善する手法で、点群を画像上に再投影したときに、観測データに最も当てはまるような3次元点群とカメラ姿勢を推定する問題です。

Visual SLAMの代表的手法であるORB-SLAM2では、ループ閉じ込みが行われたタイミングで、これまでに復元した全ての3次元点群と、全てのキーフレームの姿勢を改善するためにバンドル調整が用いられます。このバンドル調整は、論文中ではMotion Only BAやLocal BAと区別し、Full BAと呼ばれます。

Full BAは扱う問題の規模が大きく、大きな処理時間がかかることがあります。例えば、KITTIデータセットのsequences/00という比較的長いシーケンスでは、最後のループ閉じ込み後に行うFull BAで、約12秒の処理時間がかかることが分かりました。

この結果に対する「CUDAで実装したらどの程度速くなるのか」という興味・疑問が、このプロジェクトを始めたきっかけになります。うまく高速化できれば、最適化開始から点群と姿勢を更新するまでの遅延を短縮したり、CPUのリソースを軽減するといったメリットがあります。

バンドル調整のアルゴリズム

問題の定式化

$i$番目のカメラの姿勢を${\bf x}_i^{\rm p} \in \mathbb{R}^6$、$j$番目の3次元点の座標を${\bf x}_i^{\rm l} \in \mathbb{R}^3$とし、これらを並べたものを${\bf x}$とします。${\bf x}$がバンドル調整における未知数となります。${\bf x}$の次元$N$は、カメラ数を$N_{\rm p}$、3次元点数を$N_{\rm l}$として$N=6N_{\rm p}+3N_{\rm l}$です。

$$
{\bf x} = (
{\bf x}_1^{\rm pT},\cdots,{\bf x}_i^{\rm pT},\cdots,{\bf x}_{N_{\rm p}}^{\rm pT}
\ | \
{\bf x}_1^{\rm lT},\cdots,{\bf x}_j^{\rm lT},\cdots,{\bf x}_{N_{\rm l}}^{\rm lT}
) ^{\rm T} \in \mathbb{R}^N
$$

$i$番目のカメラで撮影した画像中の、$j$番目の3次元点に対応する特徴点を観測と呼び${\bf z}_{ij}$で表します。また、$i$番目のカメラに$j$番目の3次元点を再投影した点と、${\bf z}_{ij}$との誤差を再投影誤差と呼び、${\bf e}({\bf x})_{ij}$で表します。

$$
{\bf e}({\bf x})_{ij} = {\bf z}_{ij} – {\rm project}({\bf x}_i^{\rm p},{\bf x}_j^{\rm l})
$$

バンドル調整は再投影誤差の二乗和$F({\bf x})$を目的関数とし、$F({\bf x})$を最小化する${\bf x}^*$を求める問題として定式化されます。

$$
\begin{array}{l}
F({\bf x}) &=& \sum_{ij} {\bf e}({\bf x})_{ij}^{\rm T} {\bf \Omega}_{ij} {\bf e}({\bf x})_{ij} \\
{\bf x}^* &=& {\mathop{\rm arg~min}\limits_{{\bf x}}} F({\bf x})
\end{array}
$$

${\bf \Omega}_{ij}$は重み行列で、観測${\bf z}_{ij}$の不確かさ(分散)から計算されます。

非線形最小二乗法

バンドル調整の問題は非線形最小二乗問題に分類され、解法としてGauss-Newton法やLevenberg-Marquardt 法が良く利用されます。Gauss-Newton法のアルゴリズムは以下の通りです。

• 解${\bf x}$を初期化
  • ${\bf x} \leftarrow {\bf x}_0$
• 終了条件を満たすまで以下を繰り返し
  1. 現在の${\bf x}$における再投影誤差${\bf e}_{ij}$、ヤコビアン${\bf J}_{ij}$を求める
  2. ${\bf H}=\sum_{ij} {\bf J}_{ij}^{\rm T}{\bf \Omega}_{ij}{\bf J}_{ij},{\bf b}=\sum_{ij} {\bf J}_{ij}^{\rm T}{\bf \Omega}_{ij}{\bf e}_{ij},$を計算
  3. 線形方程式${\bf H} \Delta{\bf x} = -{\bf b}$を解く
  4. 解の更新：${\bf x} \leftarrow {\bf x} + \Delta{\bf x}$

線形方程式において$({\bf H}+\lambda{\bf I}) \Delta{\bf x} = -{\bf b}$としたものがLevenberg-Marquardt法(LM法)です。$\lambda$は更新量$\Delta{\bf x}$の大きさを調節するパラメータで、dumping factorと呼ばれます。

係数行列の構造

線形方程式${\bf H} \Delta{\bf x} = -{\bf b}$において、${\bf H} \in \mathbb{R}^{N \times N}$のサイズは$N=6N_{\rm p}+3N_{\rm l}$と、カメラ数と点群数に応じて大きくなります。冒頭で紹介した12秒かかるケースでは、$N$が約40万と大規模になり、直接解くことが厳しくなってきます。

そこでバンドル調整では${\bf H}$の構造を利用して、方程式をより小規模な問題に分割して解きます。${\bf H} \Delta{\bf x} = -{\bf b}$は以下の区分行列の形で表現することができます。

$$
\left(\begin{array}{cc}
{\bf H}_{\rm pp} & {\bf H}_{\rm pl} \\
{\bf H}_{\rm pl}^{\rm T} & {\bf H}_{\rm ll}
\end{array}\right)
\left(\begin{array}{c}
\Delta{\bf x}_{\rm p} \\
\Delta{\bf x}_{\rm l}
\end{array}\right)
=\left(\begin{array}{c}
-{\bf b}_{\rm p} \\
-{\bf b}_{\rm l}
\end{array}\right)
$$

${\bf H}_{\rm pp},{\bf H}_{\rm pl},{\bf H}_{\rm ll}$はそれぞれブロック疎行列(特定のブロック以外は$\bf 0$)という特徴があります。特に、${\bf H}_{\rm pp},{\bf H}_{\rm ll}$は対角ブロック以外が$\bf 0$の、ブロック対角行列です。

さて、$\left(\begin{array}{cc} {\bf A} & {\bf B} \\ {\bf C} & {\bf D}\end{array}\right)$という形の区分行列に対して、シューア補行列(Schur complement)という行列が定義され、${\bf H}$の${\bf H}_{\rm ll}$に対するシューア補行列${\bf H}_{\rm sc}$は以下のように定義されます。

$$
{\bf H}_{\rm sc} = {\bf H}_{\rm pp} − {\bf H}_{\rm pl} {\bf H}_{\rm ll}^{-1} {\bf H}_{\rm pl}^{\rm T}
$$

記載は省略しますが、シューア補行列を用いて${\bf H}$の逆行列${\bf H}^{-1}$を表現することができ、$\Delta{\bf x}=-{\bf H}^{-1}{\bf b}$を整理すると、$\Delta{\bf x}_{\rm p},\Delta{\bf x}_{\rm l}$に関する2つの方程式を得ます。

$$
\begin{array}{l}
{\rm 1}. \
{\bf H}_{\rm sc} \Delta{\bf x}_{\rm p} &=&
-{\bf b}_{\rm p} + {\bf H}_{\rm pl} {\bf H}_{\rm ll}^{-1} {\bf b}_{\rm l} \\
{\rm 2}. \
{\bf H}_{\rm ll}\Delta{\bf x}_{\rm l} &=& -{\bf b}_{\rm l} – {\bf H}_{\rm pl}^{\rm T} \Delta{\bf x}_{\rm p}
\end{array}
$$

上記より、まず$1.$を解いて$\Delta{\bf x}_{\rm p}$が求められ、次に$\Delta{\bf x}_{\rm p}$を用いて$2.$を解いて$\Delta{\bf x}_{\rm l}$が求められます。式中の${\bf H}_{\rm ll}^{-1}$は、${\bf H}_{\rm ll}$がブロック対角行列であるため、計算が容易という性質があります。$2.$を解く際も${\bf H}_{\rm ll}^{-1}$を直接利用できます。

実質的に解く必要がある線形方程式は$1.$のみになり、問題の規模も$6N_{\rm p}$に縮小されます。${\bf H}_{\rm sc}$は正定値対称行列という性質があり、コレスキー分解を用いて${\bf H}_{\rm sc}={\bf L}{\bf L}^{\rm T}$と分解することで線形方程式を解くことができます。

CUDA化の内容

リファレンス実装

CUDA化にあたり、ORB-SLAM2で使用されているグラフ最適化ライブラリg2oをリファレンス実装としました。g2oの内、バンドル調整で行われる処理の抽出、各処理のCUDA化の検討・実装を行いました。

結論から言うと、ほとんどの処理は並列化と相性が良く、CUDA化によって高速化効果が得られました。もともとの実装規模が大きく、全ての処理をCUDA化するのには手間がかかりましたが…

再投影誤差と係数行列の計算

現在の解における再投影誤差${\bf e}_{ij}$を計算し、続いてヤコビアン${\bf J}_{ij}$、係数行列${\bf H}=\sum_{ij} {\bf J}_{ij}^{\rm T}{\bf \Omega}_{ij}{\bf J}_{ij},{\bf b}=\sum_{ij} {\bf J}_{ij}^{\rm T}{\bf \Omega}_{ij}{\bf e}_{ij},$を計算します。これらの計算量は与えられた$ij$の組み合わせ、すなわち観測の数に比例します。個々の要素は累積計算を除いて独立に計算できるため、並列化による高速化効果が期待できます。

シューア補行列の計算

$\Delta{\bf x}_{\rm p}$を求めるにあたり、まず${\bf H}_{\rm sc} = {\bf H}_{\rm pp} − {\bf H}_{\rm pl} {\bf H}_{\rm ll}^{-1} {\bf H}_{\rm pl}^{\rm T}$と${\bf b}_{\rm sc}=-{\bf b}_{\rm p} + {\bf H}_{\rm pl} {\bf H}_{\rm ll}^{-1} {\bf b}_{\rm l}$を計算します。ここではブロック疎行列の演算がメインになります。計算に現れる各行列の非ゼロブロックは特殊なパターンをもち、かつ反復中に変化しないため、これに特化したCUDA実装を行っています。

${\bf H}_{\rm ll}^{-1}$は${\bf H}_{\rm ll}$がブロック対角行列のため、計算が容易です。各ブロックで$3\times3$の逆行列を計算します。次に${\bf H}_{\rm pl} {\bf H}_{\rm ll}^{-1}$を計算します。疎行列同士の積ですが、片方(${\bf H}_{\rm ll}^{-1}$)がブロック対角行列のため、実装がシンプルで済みます。

続いて${\bf H}_{\rm pl} {\bf H}_{\rm ll}^{-1}$を用いて$({\bf H}_{\rm pl} {\bf H}_{\rm ll}^{-1})\times{\bf b}_{\rm l}$と$({\bf H}_{\rm pl} {\bf H}_{\rm ll}^{-1})\times{\bf H}_{\rm pl}^{\rm T}$を計算します。前者は疎行列ベクトル積で、比較的実装が容易です。後者は疎行列同士の積で、リファレンス実装を素直にCUDA化するとスレッド毎に処理負荷がばらつき、難しい個所でした。ここでは、行列積で発生するブロック同士の乗算を単位として並列化し、スレッド数を稼ぎつつ処理負荷を均等化する工夫をしました。

線形方程式の求解

線形方程式${\bf H}_{\rm sc} \Delta{\bf x}_{\rm p} =-{\bf b}_{\rm sc}$を解きます。前述の通り、ここでは(疎行列の)コレスキー分解を利用します。g2oでは、CSParseやCHOLMOD、Eigen等、ライブラリを選択できるようになっています。本プロジェクトではCUDA Toolkit付属のcuSPARSEとcuSOLVERを利用しました。cuSPARSEやcuSOLVERはこれまで注目していませんでしたが、バンドル調整に利用できるのは嬉しいですね。

$\Delta{\bf x}_{\rm p}$が求まったら、$\Delta{\bf x}_{\rm l} = {\bf H}_{\rm ll}^{-1} (-{\bf b}_{\rm l} – {\bf H}_{\rm pl}^{\rm T} \Delta{\bf x}_{\rm p})$を計算します。疎行列ベクトル積がメインの処理です。

解の更新

最後に、$\Delta{\bf x}_{\rm p},\Delta{\bf x}_{\rm l}$を用いて解を更新します。これも各カメラ姿勢と各3次元点毎に並列化できます。

高速化結果

ORB-SLAM2とKITTIデータセットで取得したカメラ姿勢と3次元点群を入力として、処理時間をリファレンス実装のg2oと比較しました。計測プログラムと入力データは公開したソースコードに同封しているため、再現可能です。

高速化結果は以下の通りです。CPUで約12秒程度の処理をGPUで約1.2秒まで短縮できました。

計測環境

処理時間

処理時間内訳

GPUで1.2[sec]のケースにおける処理時間内訳は以下の通りで、処理時間のほとんどはコレスキー分解でした。さらなる高速化を目指す場合は、現状cuSOLVERに頼っているコレスキー分解を何とかする必要がありそうです。

ソースコードの公開

冒頭で紹介した通り、今回のソースコードはcuda-bundle-adjustmentという名前でGithubに公開しました。ORB-SLAM2内で使用する方法も紹介していますので、興味がある方は是非使ってみてください。
fixstars/cuda-bundle-adjustment

g2oに比べると現状の機能は限定的ですが、今後需要があれば機能追加や応用範囲を広げていきたいと思います。

おわりに

CUDAによるバンドル調整の高速化についてご紹介させていただきました。全ての処理をCUDA化し、納得できる処理時間にするまでには結構な試行錯誤を重ねましたが、こうして一定の高速化効果が得られ、ソースコード公開まで至ることができ嬉しく思います。

最後に、g2oの調査や高速化にご協力いただいた、昨年アルバイトの風見亮佑さんの多大な貢献に感謝いたします。

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください。