【コード解説】Einsumによるアフィン変換の実装

作成日: 2023年10月19日(木) 00:00

最終更新日: 2023年11月02日(木) 11:24

カテゴリ: コンピュータビジョン

こんにちは．高山です．
こちらの記事で，Numpyを用いたアフィン変換の実装方法を紹介しました．
このとき，追跡点配列にアフィン変換を適用する処理を np.einsum() を用いて実装したのですが，細かい部分は説明していませんでした．

np.einsum() はアインシュタインの縮約記法に基づく計算を行う関数で，内積や外積，トレースなど様々な計算処理を統一的に扱うことができます．
np.einsum() 自体は公式のAPIドキュメントに詳細な説明が載っています．
また，アインシュタインの縮約記法に関してもWeb上に解説記事がたくさんあります．
(物理学の話の中で出てくることが多いので，戸惑うかもしれませんが(^^;))
そのため，この記事ではアフィン変換式をどのように np.einsum() に当てはめていくかに焦点を絞って説明したいと思います．

今回解説するスクリプトはGitHub上に公開しています．

1. そもそもどんな処理がやりたいのか

細かな話に入る前に，そもそもどんな処理がやりたいのかを説明します．
次の式$\eqref{eqn_affine_homo}$は，以前の記事で説明したアフィン変換の式です．

$$ \begin{equation} \label{eqn_affine_homo} \begin{pmatrix} x' \\ y' \\ 1 \end{pmatrix} = \begin{pmatrix} a & b & e \\ c & d & f \\ 0 & 0 & 1 \end{pmatrix} \begin{pmatrix} x \\ y \\ 1 \end{pmatrix} \end{equation} $$

追跡点配列を変換する，という文脈で考えると，式$\eqref{eqn_affine_homo}$は1フレーム内の1個の追跡点を変換を示しています．
追跡点配列のフレーム数を $N$，1フレーム内の追跡点数を$M$ とすると，今解きたいのは $N \times M$ 個の座標点を変換する処理になります．

2. もう少し計算過程を追ってみる

式$\eqref{eqn_affine_homo}$の計算過程をもう少し追ってみましょう．
$x$座標や$y$座標などの個々の変数の意味は取り敢えず置いておきます．
式$\eqref{eqn_affine_homo}$の各変数を，添え字付き変数で置き換えて次の式$\eqref{eqn_affine_general}$のようにします．

$$ \begin{equation} \label{eqn_affine_general} \begin{pmatrix} b_0' \\ b_1' \\ b_2' \end{pmatrix} = \begin{pmatrix} a_{00} & a_{01} & a_{02} \\ a_{10} & a_{11} & a_{12} \\ a_{20} & a_{21} & a_{22} \end{pmatrix} \begin{pmatrix} b_0 \\ b_1 \\ b_2 \end{pmatrix} \end{equation} $$

式$\eqref{eqn_affine_general}$を要素ごとに展開すると，次の式$\eqref{eqn_affine_elem}$のようになります．

$$ \begin{eqnarray} b_0' & = & a_{00} b_0 + a_{01} b_1 + a_{02} b_2 \nonumber \\ \label{eqn_affine_elem} b_1' & = & a_{10} b_0 + a_{11} b_1 + a_{12} b_2 \\ b_2' & = & a_{20} b_0 + a_{21} b_1 + a_{22} b_2 \nonumber \\ \end{eqnarray} $$

式$\eqref{eqn_affine_elem}$をよく見ると，右辺の各項で$a$の2番目の添字と$b$の添字が連動している事が分かります．
この連動している添字を基に総和を用いて書き換えると，次の式$\eqref{eqn_affine_sum0}$と書けます．

$$ \begin{eqnarray} b_0' & = & \sum_{j=0}^{2} a_{0j} b_j \nonumber \\ \label{eqn_affine_sum0} b_1' & = & \sum_{j=0}^{2} a_{1j} b_j \\ b_2' & = & \sum_{j=0}^{2} a_{2j} b_j \nonumber \\ \end{eqnarray} $$

さらに式$\eqref{eqn_affine_sum0}$を見ると，$a$の1番目の添字と$b'$の添字が連動していることが分かりますので，次の式$\eqref{eqn_affine_sum1}$のように整理できます．

$$ \begin{equation} \label{eqn_affine_sum1} b_i' = \sum_{j=0}^{2} a_{ij} b_j \end{equation} $$

添字付き変数に置き換えて展開してみると，インデクスの関係が分かりやすくなると思います．

3. 追跡点配列の場合はどうなるか

次に$N \times M$ 個の座標点からなる追跡点配列を処理する場合を考えたいと思います．
と言ってもここまでくれば話はそう難しくありません．

$n$ 番目のフレーム，$m$ 番目の追跡点の $j$ 番目の座標を $b_{nmj}$ で表すことにします．
すると，式$\eqref{eqn_affine_sum1}$は次のように表すことができます．

$$ \begin{equation} b_{inm}' = \sum_{j=0}^{2} a_{ij} b_{nmj} \end{equation} $$

式$\eqref{eqn_affine_sum1}$は追跡点配列の $(i, n, m)$ 番目の値を求めていることになります．
この処理を配列の全要素に対して行えば追跡点配列全体の変換処理ができます．

4. np.einsumの処理

式$\eqref{eqn_affine_sum1}$をよく見ると，

2個の配列 (ここでは $a$ と $b$) に共通している添字 (ここでは $j$) について総和を計算
計算結果をその他の添字で表す箇所 (ここでは $(i,n,m)$ 番目) の要素として格納
(最終的には) 上記処理を全配列に対して行う

という処理になっていることが分かります．

np.einsum() は上記のルールに基づいて総和計算を行う関数で，総和計算を行う軸 (上で言うところの添字) の指定次第で様々な計算処理を行うことが可能です．

5. 実装例

ここから先は，簡単なスクリプトを用いて np.einsum() を用いた実装例をお見せします．
これから説明するスクリプトはGitHub上に公開しています．

5.1 モジュールのインポート

まず，次のコードで計算処理用に numpy と，時間計測用に time をインポートします．

import numpy as np
import time

5.2 テスト配列の初期化

次のコードでテスト用の配列を初期化します．
処理時間の違いが分かりやすいように，入力配列は少し大きなサイズに設定しました．
また，今回は配列の中身自体には注目していないので，ランダムな値で初期化をしています．

inputs = np.random.rand(1000, 1000, 3)
mat = np.random.rand(3, 3)
print(inputs.shape)
print(mat.shape)

(1000, 1000, 3)
(3, 3)

5.3 Pythonループを用いた計算

np.einsum() の処理例の前に，全配列の計算をPythonのループ処理で実装した例を下記に示します．
7-8行目に示すように，Pythonのループ処理で1個づつ追跡点を取り出してアフィン変換を適用しています．
計算結果は所望の形状になるようにリストに積んでおき，最後に12行目で np.array に変換しています．

def simple_affine(inputs, mat):
    newinputs = []
    shape = inputs.shape
    for i in range(shape[0]):
        row = []
        for j in range(shape[1]):
            temp = inputs[i, j]
            temp = np.matmul(mat, temp)
            row.append(temp)
        newinputs.append(row)
    newinputs = np.array(newinputs)
    return newinputs

処理を実行してみます．
時間計測のために同じ処理を10回繰り返して，平均時間を出力しています．

# Apply affine.
trial = 10
start = time.perf_counter()
for _ in range(trial):
    newinputs1 = simple_affine(inputs, mat)
interval = time.perf_counter() - start
print(newinputs1.shape)
print(f"Time1:{interval / trial}")

print処理の結果を示します．
処理時間は $2.95$ 秒程度でした．
こちらの処理時間をベースラインとして，np.einsum() を用いたときに処理時間がどのように変わるかを見ていきます．

(1000, 1000, 3)
Time1:2.967733277299999

5.4 np.einsum を用いた計算

np.einsum() を用いた実装例を下記に示します．
先に示したPythonループの処理が1行だけ (3行目) で実装できることが分かります．

start = time.perf_counter()
for _ in range(trial):
    newinputs2 = np.einsum("ij,nmj->nmi", mat, inputs)
interval = time.perf_counter() - start
print(f"Time2:{interval / trial}")

diff = (newinputs1 - newinputs2).sum()
print(f"Diff:{diff}")

np.einsum() を使う上で戸惑う箇所が第1引数 (ここではij,nmj->nmi) の与え方だと思います．
第1引数の意味しているところは下記のとおりです．

【コード解説】
- `,`の前は第2引数で指定した配列 (ここでは`mat`) の各軸を示す．
  `mat` は `[3, 3]` の配列で，`i` が第1軸，`j` が第2軸になります．
  (要素数が同じなので分かりにくくて済みません(^^;))

- `,`の後は第3引数で指定した配列 (ここでは`inputs`) の各軸を示す．
  `inputs` は `[1000, 1000, 3]` の配列で，`n` が第1軸，`m` が第2軸，`j` が第3軸になります．

- `j` が共通なので，`mat` の第2軸と `inputs` の第3軸に対して総和計算を行う．

- `->` の後は出力配列の形状と，総和計算の格納先を示す．
  - 出力配列は `[inputsの第1軸長, inputsの第2軸長, matの第1軸長]` の形状になり，
  - `mat` の `i` 番目の要素 (部分配列) と，`inputs` の `(n, m)` 番目の要素 (部分配列) の
    総和計算は `(n, m, i)` 番目に格納されます．

print処理の結果を示します．
処理時間は $35$ ミリ秒程度でした．
Pythonのループ処理が無くなり，Numpyの高速な実装で処理されるようになったので処理時間が大きく改善しています．
また，処理結果の差はほぼ無いことが分かります．

Time2:0.03701798059999959
Diff:-6.526723606015139e-14

5.5 np.einsum 処理の他の書き方

np.einsum() は計算処理手順を引数で指定するので，上の例とは別の書き方もできます．
2個ほど例をお見せします．

Implicit mode

下の例では，np.einsum("abj,ij", inputs, mat) という書き方で同じ処理を行っています．
まず，第2引数と第3引数を入れ替えています．
これに伴って，第1引数の , 前後の軸指定が入れ替わっています．
また，第1引数の -> 以降を省略しています．
このような書き方を Implicit mode と言います．
(一方，-> を付けた書き方を Explicit mode と呼びます) この場合，出力形式は abj,ij から自動的に推論されます．
具体的には， ab,i をアルファベット順に並べた出力形式になります．

start = time.perf_counter()
for _ in range(trial):
    newinputs2 = np.einsum("abj,ij", inputs, mat)
interval = time.perf_counter() - start
print(f"Time3:{interval / trial}")

diff = (newinputs1 - newinputs2).sum()
print(f"Diff:{diff}")

print処理の結果を示します．
処理時間，出力配列ともに同様の結果が得られています．

Time3:0.03361054859999939
Diff:-6.526723606015139e-14

Ellipsisオブジェクトの利用

下の例では，np.einsum("...j,ij->...i", inputs, mat) という書き方で同じ処理を行っています．
第2，第3引数の入れ替えは上の例と同じです．
","の前と"->"の後ろの"..."はEllipsisオブジェクトと言って，専ら引数やインデクスなどを省略するために使用します．
ここでは上の例における n, m 軸を省略しています．

この書き方の利点は，inputs の軸数を可変にできる点です．
ここでは inputs の形状は [1000, 1000, 3] ですが，Ellipsisを用いて処理を実装することで，[1, 1000, 1000, 3] など新しい軸を加えた場合でもそのまま動作します．

start = time.perf_counter()
for _ in range(trial):
    newinputs2 = np.einsum("...j,ij->...i", inputs, mat)
interval = time.perf_counter() - start
print(f"Time4:{interval / trial}")

diff = (newinputs1 - newinputs2).sum()
print(f"Diff:{diff}")