AdaGradは学習率を自動調整してくれる勾配法の亜種で、いろんな人が絶賛しています。 勾配を足し込む時に、各次元ごとに今までの勾配の2乗和をとっておいて、その平方根で割ってあげるだけと、恐ろしくシンプルです。
Adaptive Subgradient Methods for Online Learning and Stochastic OptimizationJohn Duchi, Elad Hazan, Yoram Singer. JMLR 2011.
丁度、 @echizen_tm さんがブログを書いてました。
AdaGrad+RDAを実装しました。通常のSGDなどは学習率をだんだん減衰させながら勾配を足していくわけですが、どの様に減衰させるかという問題にいつも頭を悩ませます。 AdaGradでは最初の学習率こそ外から与えますが、減衰のさせ方や減衰率といったハイパーパラメータから開放されます。 今日の話は、そんなAdaGradがちょっとの工夫で12倍速になった話です。
さて、本題はここから。 AdaGradの部分を高速化したいなー、と思って更新部分だけ抜き出して速度を測ってみました。 更新部分を抜粋。
void ada_grad(
float eta,
const std::vector<float>& g,
std::vector<float>& sum,
std::vector<float>& x) {
for (std::size_t i = 0, size = x.size(); i < size; ++i) {
sum[i] += g[i] * g[i];
x[i] += eta / std::sqrt(sum[i]) * g[i];
}
}
コードの全体は以下のgistにあります。
https://gist.github.com/unnonouno/9ca83eefb31a6474caf6試しに速度を測ってみましょう。
$ gcc -O2 adagrad_test.cpp $ ./a.out 7744.53 msec 200.099
まぁ、こんなもんです。 下の200.099は計算結果を適当に出力してるだけ。 さて、ここで一工夫。 なんてことはなくて、コンパイルオプションを変えます。 では、もう1度。
$ gcc -Ofast -march=native adagrad_test.cpp $ ./a.out 540.679 msec 200.099
驚異的に速くなりましたね。 -O2の12倍です。 ちょっと何いってるのかわからないレベルで速くなります。 もちろんコンパイラが頑張るからなんですが、何でこんなに速くなるんでしょう。 アセンブリコードをよんでみましょう。 まずは、-O2から。 該当部分だけ。
.L5:
movss (%rcx,%rbx,4), %xmm1
movq (%rsi), %rax
mulss %xmm1, %xmm1
leaq 0(,%rbx,4), %rbp
addq %rbp, %rax
addss (%rax), %xmm1
sqrtss %xmm1, %xmm0
ucomiss %xmm0, %xmm0
movss %xmm1, (%rax)
jp .L10
.L3:
movaps %xmm3, %xmm1
movq (%rdx), %rax
addq $1, %rbx
divss %xmm0, %xmm1
addq %rbp, %rax
cmpq %r12, %rbx
mulss (%rcx,%rbp), %xmm1
addss (%rax), %xmm1
movss %xmm1, (%rax)
jne .L5
まぁこんなものかな。 そして、-Ofastの方。
.L4:
vmovups (%rdi,%rax), %xmm1
addq $1, %r8
vmovups (%rsi,%rax), %xmm2
vinsertf128 $0x1, 16(%rdi,%rax), %ymm1, %ymm1
vmulps %ymm1, %ymm1, %ymm1
vinsertf128 $0x1, 16(%rsi,%rax), %ymm2, %ymm2
vaddps %ymm2, %ymm1, %ymm1
vmovups %xmm1, (%rsi,%rax)
vextractf128 $0x1, %ymm1, 16(%rsi,%rax)
vrsqrtps %ymm1, %ymm2
vmulps %ymm1, %ymm2, %ymm1
vmulps %ymm2, %ymm1, %ymm1
vaddps %ymm5, %ymm1, %ymm1
vmovups (%rcx,%rax), %xmm3
vmulps %ymm4, %ymm2, %ymm2
vmulps %ymm2, %ymm1, %ymm2
vmovups (%rdi,%rax), %xmm1
vinsertf128 $0x1, 16(%rcx,%rax), %ymm3, %ymm3
vmulps %ymm6, %ymm2, %ymm2
vinsertf128 $0x1, 16(%rdi,%rax), %ymm1, %ymm1
vmulps %ymm1, %ymm2, %ymm1
vaddps %ymm3, %ymm1, %ymm1
vmovups %xmm1, (%rcx,%rax)
vextractf128 $0x1, %ymm1, 16(%rcx,%rax)
addq $32, %rax
cmpq %r8, %r10
ja .L4
劇的に変わりました。 まず、全体的にvのつく命令に置き換わっています。 そう、AVX命令に置き換わってます。 AVXでは、8つの単精度浮動少数を1度に扱えるので、これで8倍速です。 ループ内は掛けたり割ったりしてますが、分岐がないのでベクトル化できちゃう。 すごい。
でももっと速くなってますね。 何ででしょう。 ボトルネックになりそうなsqrt周りを見てみましょう。 実は-O2でも、sqrtssというCPU命令を使って1発で計算しているのでそれなりに高速です。 ところが、-Ofastでは、vrsqrtpsという何やら怪しげな命令が。 rsqrtというのは、-1/2乗を計算する命令です。 そのお陰で、-Ofast側ではdivが無くなっていることに気づきます。
ところで、何でrsqrtという命令が有るかというと、1/2乗よりも-1/2乗の方が簡単だからと思われます。 これは学生実験の時に教わった気がします。 -1/2をニュートン法で計算するときに、わり算が必要ないのです。
http://www.riken.jp/brict/Ijiri/study/fastsqrt.html加えて、 rsqrtは精度を落として計算しているため高速なようです だいたい、初期学習率が適当なので、今のケースで精度はさしていらないでしょう。 ちょっと調べてみると、sqrt 命令が8〜105サイクルなのに rsqrt 命令が驚きの1サイクルと書かれています。 ということは、sqrtよりも高効率で、しかもdivも消えちゃう。 ちなみにperfで見てみると、rsqrtは全体の5%しかCPU時間を食っておらず、その他の掛け算1回分と同じコストでした。 極めて効率的に計算できるようです。
もはやsqrtがボトルネックになる時代ではないんですね。 対象とするCPUがどのような機能を持っていて、利用するコンパイラがどのような最適化をかけてくれるのかということを把握することは大事です。 実は最初実装した時に、分母が0のときに割らないように分岐していたのですが、εを足すように変えました。 これにしないと最適化がかかりません。 そもそも、掛けたり足したりしてるし、ベクトル2つ使ってるし、sqrtもでてくるし、流石にベクトル化できないだろうと思ってたら甘かったです。 ちなみにこのとき、1回の学習に8時間掛かるという試算になっていたのですが、上記の最適化とOpenMPによる4倍速で、計50倍速、最終的に10分で終わるようになりました。 遅いなー、と思って10台のマシンで計算するよりも、コンパイラの最適化しやすいように工夫したほうが効果が高かったりするんですね。
0 件のコメント:
コメントを投稿