Vim scriptで自動微分してみた
はじめに
Vim scriptで自動微分ライブラリを書いてみました。
この記事では、そのライブラリについて、例を示しながら紹介したいと思います。
vim-autogradとは
vim-autogradは、PyTorchやChainerのようなDefine-by-Runスタイルで自動微分を行えるライブラリで、純粋なVim scriptで書かれています。自動微分自体は、様々な分野で利用されており、近年ではディープラーニングの最適化の根幹を担う技術となっています。一般的に、forwardモードとreverseモードがありますが、現段階ではreverseモードのみ対応しています。
導入方法
このライブラリはシンプルな構成をしており、一般的なVimプラグインマネージャであれば簡単に導入できます。
vim-plugの場合、次のようにインストールできます。
Plug 'pit-ray/vim-autograd'
このライブラリを使ったプラグインを作るときは、MITライセンスの範囲内でバンドルするか、ユーザにvim-autogradのインストールを促すことになります。
基本的な微分
vim-autogradで微分したい関数を定義する場合、vim-autogradが提供する微分可能な関数を利用する必要があります。
現在(2022/03)時点では、以下の関数が実装されています。
autograd#add()autograd#mul()autograd#sub()autograd#div()autograd#pow()autograd#sqrt()autograd#log()autograd#exp()autograd#sin()autograd#cos()autograd#tanh()autograd#abs()autograd#sign()autograd#sum()autograd#broadcast_to()autograd#sum_to()autograd#max()autograd#maximum()autograd#transpose()autograd#matmul()autograd#reshape()autograd#flatten()
引数や動作はnumpyやPyTorchをベースとしています。ただし、Vim scriptでは演算子をオーバーロードすることができないため、基本的な演算に関しては、短縮版メソッドを用意しています。
| 短縮形 | 呼ばれる関数 |
|---|---|
.a(other) |
autograd#add(self, other) |
.m(other) |
autograd#mul(self, other) |
.s(other) |
autograd#sub(self, other) |
.d(other) |
autograd#div(self, other) |
.p(other) |
autograd#pow(self, other) |
.n() |
autograd#mul(self, -1.0) |
例えば、
function! s:f(x) abort
let x = a:x
let y = autograd#sub(x.p(5), x.p(3).m(2))
return y
endfunction
ここで定義した関数に対し、入力
let x = autograd#tensor(2.0)
let y = s:f(x)
すると、vim-autogradは次のような計算グラフを構築します。
このグラフの末端ノードである backward()メソッドで逆伝播を開始すると、連鎖律に従って出力側から入力側に勾配が伝わり、x.gradに
call y.backward()
echo x.grad.data
output
[56.0]
Tensorオブジェクトはいくつかの属性を持ちますが、よく使うのは.dataと.gradと.shapeです。
-
.data
Tensorのもつの生のデータであり、1次元のFloatのリストであることが保証されます。仮に、autograd#tensor([[1, 2], [3, 4]])と定義しても、[1.0, 2.0, 3.0, 4.0]に変換されます。 -
.grad
ある出力のその変数に関する勾配です。つまり、次のように中間オブジェクトを作ると、途中の勾配を抽出できます。let x = autograd#tensor(2.0) let t = autograd#mul(x, 3.0) let y = autograd#exp(t) call y.backward() " x.grad: dy/dx " t.grad: dy/dt.gradは入れ子のTensorオブジェクトであるため、x.grad.dataのようにして生のデータにアクセスできます。 -
.shape
配列の形状を表すNumberのリストを保持します。.sizeや.dim属性はなく、次のようにlen()関数を使います。let x = autograd#tensor([[1, 2], [3, 4]]) " x.shape : [2, 2] " x.data : [1.0, 2.0, 3.0, 4.0] let dim = len(x.shape) " dim: 2 let size = len(x.data) " size: 4
全体のコード
高階微分
高階微分とは、微分の微分であり、逆伝播の計算グラフに対するバックプロパゲーションと言い換えることができます。vim-autogradは、バックプロパゲーションに対する計算グラフ生成機能(double backprop)をサポートしているため、同様の手順で高階微分を行うことができます。
ただし、vim-autogradにおける高階微分は、以下の二つの点に注意しなければなりません。
-
勾配の蓄積
vim-autogradは、Tensorオブジェクトのbackward()メソッドを呼び出した段階で.gradを確保しますが、既に.gradを保持している場合には加算を行います。これは、「同じ変数の使いまわし」の逆伝播が加算であるためです[1]。したがって、backward()メソッドを呼ぶたびに勾配をリセットしなければ、蓄積された勾配に加算された値が返ってきます。次は間違った例です。
間違った例x = autograd#tensor(2.0) y = x.p(2).a(5) " y = x^2 + 5 " ➀ 1階微分 call y.backward(1) let gx = x.grad echo gx.data " => [4.0] " ➁ 2階微分 call gx.backward() echo x.grad.data " => [6.0]関数
f(x)=x^2+5 f'(x)=2x f'(2)=4 f''(x)=2
正しい例x = autograd#tensor(2.0) y = x.p(2).a(5) " y = x^2 + 5 " ➀ 1階微分 call y.backward() let gx = x.grad echo gx.data " => [4.0] " 勾配のリセット call x.cleargrad() " ➁ 2階微分 call gx.backward() echo x.grad.data " => [2.0]このように、高階微分では
backward()の直前にcleargrad()メソッドで勾配のリセットを行う必要があります。ただ、このような手続きは面倒なので、PyTorchのように勾配を蓄積しないautograd#grad()関数を提供しています。この関数を用いることで、引数に渡したTensorオブジェクトの.grad属性を汚染することなく、微分を行うことができます。次の例は全く同じことを行っています。
.backward()let y = f(x) call y.backward(1) let gx1 = x.grad call x.cleargrad() call gx1.backward() let gx2 = x.gradautograd#grad()
let y = f(x) let gx1 = autograd#grad(y, x, 1) let gx2 = autograd#grad(gx1, x, 1)以上のように、その後にさらなる微分が控えている場合には、
autograd#grad()関数で微分するか、backward()メソッドの直前にcleargrad()を呼びます。 -
バックプロパゲーションのモード
vim-autogradは、引数なしでbackward()メソッドやautograd#grad()関数を呼んだ場合、バックプロパゲーション自体に対する計算グラフを生成しません。したがって、高階微分を行う際には、次のcreate_graphフラグをTrueにする必要があります。Tensor.backward(create_graph=0, retain_outgrad=0)autograd#grad(output, inputs, create_graph=0, retain_outgrad=0)
以上に注意して、関数
let x = autograd#tensor(2.0)
" y = x^5 - 2x^3 + 4x^2 + 6x + 5
let y = s:f(x)
echo 'y :' y.data
" gx1 = 5x^4 - 6x^2 + 8x + 6
let gx1 = autograd#grad(y, x, 1)
echo 'gx1:' gx1.data
" gx2 = 20x^3 - 12x + 8
let gx2 = autograd#grad(gx1, x, 1)
echo 'gx2:' gx2.data
" gx3 = 60x^2 - 12
call gx2.backward(1)
echo 'gx3:' x.grad.data
output
y : [49.0]
gx1: [78.0]
gx2: [144.0]
gx3: [228.0]
全体のコード
ディープラーニングを利用したワイン分類
自動微分を用いると、損失関数の値が最小となるように、パラメータを最適化できます(勾配降下法)。ここではUCIで提供されているアルコールや酸の量などの13次元のベクトルから3種類のワインを分類するToyデータセットを利用して、ディープラーニングを行っていきます。データセットの詳細は以下の通りです。
出典: UCI Machine Learning Repository: Wine Data Set
データセットの前処理
まず初めに、データセットを次のような形式の辞書として返す関数をs:get_wine_dataset()関数として実装します。
{
'train': [
[[13ベクトル], [13ベクトル], ...],
[[正解のクラスID], [正解のクラスID], ...],
],
'test': [
[[13ベクトル], [13ベクトル], ...],
[[正解のクラスID], [正解のクラスID], ...],
],
'insize': 13,
'nclass': 3
}
この関数は、データのスケールがそれぞれ異なることを考慮し、平均0、分散1になるように標準化を併せて行います。今回は実装の詳細を省きますが、興味のある方は全体のコードのs:get_wine_dataset()関数を参照してみてください。
デフォルトでは、訓練データが148、テストデータが30個に分割されます。
ニューラルネットワークの構築
次に、vim-autogradが提供する関数を用いて、ニューラルネットワークを構築します。
function! s:linear(x, W, b={}) abort
let t = autograd#matmul(a:x, a:W)
return empty(a:b) ? t : autograd#add(t, a:b)
endfunction
function! s:relu(x) abort
return autograd#maximum(a:x, 0.0)
endfunction
function! s:softmax(x) abort
let y = autograd#exp(a:x.s(autograd#max(a:x)))
let s = autograd#sum(y, 1, 1)
return autograd#div(y, s)
endfunction
function! s:cross_entropy_loss(y, t)
let loss = autograd#mul(a:t, autograd#log(a:y))
let batch_size = loss.shape[0]
return autograd#div(autograd#sum(loss), batch_size).n()
endfunction
let s:MLP = {'params': []}
function! s:MLP(in_size, ...) abort
let l:mlp = deepcopy(s:MLP)
let std = sqrt(2.0 / a:in_size)
let l:W = autograd#normal(0, std, [a:in_size, a:1])
let l:b = autograd#zeros([a:1])
let l:W.name = 'W0'
let l:b.name = 'b0'
let l:mlp.params += [l:W, l:b]
for l:i in range(a:0 - 1)
let std = sqrt(2.0 / a:000[l:i])
let l:W = autograd#normal(0, std, [a:000[l:i], a:000[l:i + 1]])
let l:W.name = 'W' . string(l:i + 1)
let l:b = autograd#zeros([a:000[l:i + 1]])
let l:b.name = 'b' . string(l:i + 1)
let l:mlp.params += [l:W, l:b]
endfor
return l:mlp
endfunction
function! s:MLP.forward(x) abort
let y = s:linear(a:x, self.params[0], self.params[1])
for l:i in range(2, len(self.params) - 1, 2)
let y = s:relu(y)
let y = s:linear(y, self.params[l:i], self.params[l:i + 1])
endfor
let y = s:softmax(y)
return y
endfunction
今回定義したのは、活性化関数がReLU、出力の関数がSoftmaxの全結合層からなるニューラルネットワークです[2]。
SGDの実装
次に、Momentum、Weight Decay、Gradient Clippingを備えたSGDの実装を示します[3]。
let s:SGD = {
\ 'vs': {},
\ 'momentum': 0.9,
\ 'lr': 0.01,
\ 'weight_decay': 0.0,
\ 'grad_clip': -1
\ }
function! s:SGD.each_update(param) abort
if self.weight_decay != 0
call autograd#elementwise(
\ [a:param.grad, a:param],
\ {g, p -> g + self.weight_decay * p}, a:param.grad)
endif
if self.momentum == 0
return autograd#elementwise(
\ [a:param, a:param.grad], {p, g -> p - g * self.lr}, a:param)
endif
if !self.vs->has_key(a:param.id)
let self.vs[a:param.id] = autograd#zeros_like(a:param)
endif
let v = self.vs[a:param.id]
let v = autograd#sub(v.m(self.momentum), a:param.grad.m(self.lr))
let self.vs[a:param.id] = v
return autograd#elementwise([a:param, v], {a, b -> a + b}, a:param)
endfunction
function! s:SGD.step(params) abort
" gradients clipping
if self.grad_clip > 0
let grads_norm = 0.0
for param in a:params
let grads_norm = autograd#sum(param.grad.p(2))
endfor
let grads_norm = autograd#sqrt(grads_norm).data[0]
let clip_rate = self.grad_clip / (grads_norm + 0.000001)
if clip_rate < 1.0
for param in a:params
let param.grad = param.grad.m(clip_rate)
endfor
endif
endif
call map(a:params, 'self.each_update(v:val)')
endfunction
function! s:SGD(lr=0.01, momentum=0.9, weight_decay=0.0, grad_clip=-1) abort
let l:optim = deepcopy(s:SGD)
let l:optim.lr = a:lr
let l:optim.momentum = a:momentum
let l:optim.weight_decay = a:weight_decay
let l:optim.grad_clip = a:grad_clip
return l:optim
endfunction
学習と評価
上で定義したネットワークとオプティマイザを利用することで、簡単に学習させられます。まず、データをランダムにサンプリングし、ミニバッチを作ります。その後、順伝播と逆伝播を行ったあとに、SGDで最適化を行います。最後に、訓練済みのモデルに対し、テストデータをもとにした単純な正解率を算出します。
function! s:main() abort
call autograd#manual_seed(42)
let data = s:get_wine_dataset()
let model = s:MLP(data['insize'], 100, data['nclass'])
let optimizer = s:SGD(0.1, 0.9, 0.0001, 10.0)
" train
let max_epoch = 50
let batch_size = 16
let train_data_num = len(data['train'][0])
let each_iteration = float2nr(ceil(1.0 * train_data_num / batch_size))
let logs = []
for epoch in range(max_epoch)
let indices = autograd#shuffle(range(train_data_num))
let epoch_loss = 0
for l:i in range(each_iteration)
let x = []
let t = []
for index in indices[l:i * batch_size:(l:i + 1) * batch_size - 1]
call add(x, data['train'][0][index])
let onehot = repeat([0.0], data['nclass'])
let onehot[float2nr(data['train'][1][index][0])] = 1.0
call add(t, onehot)
endfor
let y = model.forward(x)
let loss = s:cross_entropy_loss(y, t)
for param in model.params
call param.cleargrad()
endfor
call loss.backward()
call optimizer.step(model.params)
let l:epoch_loss += loss.data[0]
endfor
let l:epoch_loss /= each_iteration
" logging
call add(logs, epoch . ', ' . l:epoch_loss)
call writefile(logs, '.autograd/train.log')
endfor
" evaluate
let ng = autograd#no_grad()
let accuracy = 0.0
for l:i in range(len(data['test'][0]))
let pred = model.forward([data['test'][0][l:i]])
" argmax
let class_idx = index(pred.data, autograd#max(pred).data[0])
let accuracy += class_idx == data['test'][1][l:i][0]
endfor
call ng.end()
echomsg 'accuracy: ' . accuracy / len(data['test'][1])
endfunction
このVim scriptを実行すると、次のようにロスが減少し、数分で評価結果が出力されます。
.autograd/train.log
0, 0.379945
1, 0.094833
2, 0.029978
3, 0.002876
4, 0.027007
5, 0.065495
6, 0.020479
7, 0.01342
8, 0.046886
9, 0.042945
...
output
accuracy: 0.966667
今回構築したニューラルネットワーク(損失関数含む)の計算グラフは、次のようになります。
.autograd/loss.png
以上のように、PyTorchライクな記述でディープラーニングを行うことができました。
全体のコード
まとめと今後
Vim scriptで自動微分ライブラリを書いてみました。このライブラリを用いることで、高階微分やディープラーニングを行うことができます。今後、Vim9 scriptに対応することで、もう少し実用的なものになり、小さなニューラルネットワークを利用したプラグインなどを作成できるようになるかもしれません。
追記
Vim9 scriptに対応しました。(2023/08/21)
推論では9倍程度、学習では7倍程度、速くなっています。
参考
本ライブラリのコア設計は、『ゼロから作る Deep Learning ❸』(O'Reilly Japan, 2020)のDeZeroをベースとしました。そのほかにも、逆伝播の実装にChainer、インターフェースの設計にPyTorchやNumPyを参考にしました。
また、ニューラルネットワークの学習をVim scriptで行う点に関連して、mattnさんのvim-brainを参考にしています。
-
f(x)=x f'(x)=1 g(x)=x+x g'(x)=f'(x)+f'(x) g'(x)=2 -
vim-autogradが提供するFunctionを呼び出して合成関数の実装を行うより、逆伝播を個別に実装する方が効率が良いことが知られています。vim-autogradは、そのような利用にも対応しており、
:help autograd#Function()に詳細があります。 ↩︎ -
当初、単純なSGDではうまく収束せず、あれやこれやと機能を付け足した結果、複雑なものができてしまいました。実装を見直したら、単純なSGDでも収束しました... ↩︎
Discussion