CNNの高速化: Winograd’s Minimal Filtering

GPUの主な用途として行列積のほかに畳み込みが挙げられるようになって数年が経ちました。今日も弊社オフィスのそこかしこでGPUたちが必死に畳み込みと行列積を計算し続けています。

畳み込みの計算が速くなると、畳み込みニューラルネットワーク (CNN) を用いたモデルの学習や推論に要する時間が短くなって嬉しいわけですが、Winograd’s Minimal Filtering と呼ばれるアルゴリズムを用いると、カーネルサイズの小さい畳み込み層の計算を高速化することができ、うまく実装するとナイーブなアルゴリズムの理論ピークを越える性能をたたき出すことができるようです。

すでにcuDNNにも実装されているようなので、自分で実装する機会はあまりないかとは思います。しかし、いつディープラーニング用ライブラリが整備されていない環境でCNNを実装することになるかわからないので、その時に備えてどこが重要かくらいは理解しておきたいと思います。

以下、Fast Algorithms for Convolutional Neural Networks の該当部分 (4章前半あたり) の概要になります。

• 入力とフィルタをあらかじめ変換し、演算後に逆変換することで計算量を減らすことができる。
• 畳み込み層の計算においては、変換結果を使いまわすことができるので相対的に変換のコストが小さくなる。
• 変換先の空間での計算も行列積に帰着できるため、計算量を削りつつも効率よく処理できる。

Winograd’s Minimal Filtering Algorithm では、入力とフィルタを元の空間における畳み込みが要素ごとの積となるような空間に変換し、要素ごとの積をとった後に逆変換することで畳み込みを行います。

このアルゴリズムを用いたとき、出力サンプル数 $m$、フィルタサイズ $r$ の畳み込み $F(m, r)$ に必要となる乗算回数 $\mu(F(m, r))$ は $m + r – 1$ となります。

また、変換と逆変換をネストすることによって2次元の畳み込みも行うことができます。その場合、出力サンプル数 $m \times n$、フィルタサイズ $r \times s$ の畳み込み $F(m \times n, r \times s)$ に必要な乗算回数 $\mu(F(m \times n, r \times s))$ は $(m + r – 1)(n + s – 1)$ となります。

まず、出力サンプル数2、フィルタサイズ2の1次元の畳み込みについて考えてみます。ここで入力 $d$ とフィルタ $g$ を以下のようにおきます。

$$
d = \left[
\begin{array}{rrr}
d_0 & d_1 & d_2
\end{array}
\right]^{\mathrm{T}}
\\
g = \left[
\begin{array}{rr}
g_0 & g_1
\end{array}
\right]^{\mathrm{T}}
$$

このとき、出力 $Y$ を以下のように求めると、合わせて4回の乗算が必要となります。

$$
Y_0 = d_0 g_0 + d_1 g_1 \\
Y_1 = d_1 g_0 + d_2 g_1
$$

しかし、以下のように $m_1, m_2, m_3$ をおき、それを用いて $Y$ を求めると乗算の回数が3回となります。

$$
m_1 = (d_0 – d_1) g_0 \\
m_2 = d_1 (g_0 + g_1) \\
m_3 = (d_2 – d_1) g_1 \\
Y_0 = m_1 + m_2 \\
Y_1 = m_2 + m_3
$$

上の式における変換を行列 $A, B, G$ を用いて表すと、畳み込みは以下の式のようにあらわされます。（ここで、$\odot$ は要素ごとの積を表します。）

$$
A^{\mathrm{T}} = \left[
\begin{array}{rrr}
1 & 1 & 0 \\
0 & 1 & 1
\end{array}
\right]
\\
G = \left[
\begin{array}{rr}
1 & 0 \\
1 & 1 \\
0 & 1
\end{array}
\right]
\\
B^{\mathrm{T}} = \left[
\begin{array}{rrr}
1 & -1 & 0 \\
0 & 1 & 0 \\
0 & -1 & 1
\end{array}
\right]
\\
Y = A^{\mathrm{T}} [G g \odot B^{\mathrm{T}} d]
$$

ここまでに出てきた方法では小さいサイズの入力を処理することのみを考えていました。$m, n$ が大きい場合にも同様のアルゴリズムを適用することはできるのですが、変換・逆変換にかかるコストが大きくなってしまいます。

元の画像を小さく分割して処理し、その結果を後でまとめると、小さい入力にのみ対応した畳み込みで大きな画像を処理することができます。以下の図に、$F(2 \times 2, 3 \times 3)$ の畳み込みを用いて大きい画像を畳み込むときのイメージを示します。

CNNで使われる畳み込みの特徴として、1つの入力をフィルタを変えつつ何度も畳み込むという点が挙げられます。この性質は上述のアルゴリズムと相性がよく、畳み込みを $CK$ 回行うのに対して、入力の変換（$B^{\mathrm{T}} d B$ の計算）は $C$ 回だけ行えばよいため、出力チャンネル数が十分に多ければ変換コストの元を取りやすくなります。

また、逆変換についても同様のことをいうことができます。以下のように式を変形してやると逆変換の回数は $K$ 回となり、入力チャンネル数が十分に多ければ逆変換のコストが相対的に小さくなっていきます。

$$
\begin{aligned}
Y_{i,k,\tilde{x},\tilde{y}} & = \sum_{c=1}^C A^{\mathrm{T}} [U_{k,c} \odot V_{c,i,\tilde{x},\tilde{y}}] A \\
& = A^{\mathrm{T}} \Bigl[ \sum_{c=1}^C U_{k,c} \odot V_{c,i,\tilde{x},\tilde{y}} \Bigr] A
\end{aligned}
$$

$i, \tilde{x}, \tilde{y}$ をまとめて1つの次元 $b$ とし、$M_{k,b} = \sum U_{k,c} \odot V_{c,b}$ とおきます。また、行列 $A$ の $\xi$ 行 $\nu$ 列にある要素を $A^{(\xi,\nu)}$ と表すとすると、

$$
M_{k,b}^{(\xi,\nu)} = \sum_{c=1}^C U_{k,c}^{(\xi,\nu)} V_{c,b}^{(\xi,\nu)} \\
M^{(\xi,\nu)} = U^{(\xi,\nu)} V^{(\xi,\nu)}
$$

となり、変換後の空間における演算は行列積となることがわかります。行列積であればハードウェアの限界近くまで性能を引き出すことができるため、余計なオーバーヘッドで性能向上が相殺されるという心配もなさそうです。

おわりに

Tags