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

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)$ となります。

$F(2, 2)$

まず、出力サンプル数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]
$$

$F(2 \times 2, 3 \times 3)$

次に、出力サンプル数2、フィルタサイズ3の畳み込みについて考えてみましょう。ここで用いる入力と変換行列および畳み込みの結果はそれぞれ以下のようになります。(変換行列は著者が公開しているツールなどで求めることができます。)

$$
d = \left[
\begin{array}{rrrr}
d_0 & d_1 & d_2 & d_3
\end{array}
\right]^{\mathrm{T}}
\\
g = \left[
\begin{array}{rrr}
g_0 & g_1 & g_2
\end{array}
\right]^{\mathrm{T}}
\\
A^{\mathrm{T}} = \left[
\begin{array}{rrrr}
1 & 1 & 1 & 0 \\
0 & 1 & -1 & 1
\end{array}
\right]
\\
G = \left[
\begin{array}{rrr}
1 & 0 & 0 \\
\frac 1 2 & \frac 1 2 & \frac 1 2 \\
\frac 1 2 & -\frac 1 2 & \frac 1 2 \\
0 & 0 & 1
\end{array}
\right]
\\
B^{\mathrm{T}} = \left[
\begin{array}{rrrr}
1 & 0 & -1 & 0 \\
0 & 1 & 1 & 0 \\
0 & -1 & 1 & 0 \\
0 & -1 & 0 & 1
\end{array}
\right]
\\
Y = A^{\mathrm{T}} [G g \odot B^{\mathrm{T}} d]
$$

また、以下のように入力を行列として変換をネストすると、2次元の畳み込みを求めることができます。

$$
d = \left[
\begin{array}{rrrr}
d_{00} & d_{01} & d_{02} & d_{03} \\
d_{10} & d_{11} & d_{12} & d_{13} \\
d_{20} & d_{21} & d_{22} & d_{23} \\
d_{30} & d_{31} & d_{32} & d_{33}
\end{array}
\right]
\\
g = \left[
\begin{array}{rrr}
g_{00} & g_{01} & g_{02} \\
g_{10} & g_{11} & g_{12} \\
g_{20} & g_{21} & g_{22}
\end{array}
\right]
\\
Y = A^{\mathrm{T}} [G g G^{\mathrm{T}} \odot B^{\mathrm{T}} d B] A
$$

ナイーブな方法で $F(2 \times 2, 3 \times 3)$ の畳み込みを行った場合、$2 \times 2 \times 3 \times 3 = 36$ 回の乗算が必要となります。一方で、要素ごとの積となるように変換した場合は $(2 + 3 – 1) \times (2 + 3 – 1) = 16$ 回の乗算で済んでおり、2.25倍の差がついています。

畳み込みニューラルネットワークへの適用

タイリング

ここまでに出てきた方法では小さいサイズの入力を処理することのみを考えていました。$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)}
$$

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

おわりに

もともとのアルゴリズムは1980年頃には知られていたようなのですが、単純な画像1枚に対するフィルタ処理では変換のコストが重いこともあってか、汎用プロセッサによる畳み込み計算ではあまり用いられていませんでした。しかし、問題設定が少し変わることでそのデメリットが小さくなり、理論性能・実効性能ともに従来のアルゴリズムを超えられるということを示され、問題全体の性質を利用することや、既存手法をうまく適用できないかを注意深く検討することの大切さを思い知らされるような内容でした。

また、今回紹介したアルゴリズムを自分で実装もしてみたのですが、パラメータ次第では比較的簡単にナイーブ実装の理論ピークを超えることができました。まだ試行錯誤しているところなのでちょっとコードを公開できる状態ではないのですが、近いうちにどこかで公開することになるかと思います。

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です