Vim scriptで自動微分してみた

75674 ワード
Vim 機械学習 tech テキストリンク

はじめに

Vim scriptで自動微分ライブラリを書いてみました。
この記事では、そのライブラリについて、例を示しながら紹介したいと思います。

vim-autogradとは

vim-autogradは、PyTorchやChainerのようなDefine-by-Runスタイルで自動微分を行えるライブラリで、純粋なVim scriptで書かれています。自動微分自体は、様々な分野で利用されており、近年ではディープラーニングの最適化の根幹を担う技術となっています。一般的に、forwardモードとreverseモードがありますが、現段階ではreverseモードのみ対応しています。

https://github.com/pit-ray/vim-autograd

導入方法

このライブラリはシンプルな構成をしており、一般的な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)

例えば、y=f(x)=x^{5}-2x^{3}という関数であれば、次のように実装できます。

function! s:f(x) abort
  let x = a:x
  let y = autograd#sub(x.p(5), x.p(3).m(2))
  return y
endfunction

ここで定義した関数に対し、入力 x=2.0 を次のように流し込んでみます。

  let x = autograd#tensor(2.0)
  let y = s:f(x)

すると、vim-autogradは次のような計算グラフを構築します。

このグラフの末端ノードである y を起点にbackward()メソッドで逆伝播を開始すると、連鎖律に従って出力側から入力側に勾配が伝わり、x.gradに\frac{{\partial}y}{{\partial}x}の値が格納されます。

  call y.backward()
  echo x.grad.data

output

[56.0]

f(x)=x^{5}-2x^{3}のxに関する微分は、f'(x)=5x^{4}-6x^{2}なので、f'(2)=56と一致します。

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

全体のコード

https://github.com/pit-ray/vim-autograd/blob/main/examples/basic.vim

高階微分

高階微分とは、微分の微分であり、逆伝播の計算グラフに対するバックプロパゲーションと言い換えることができます。vim-autogradは、バックプロパゲーションに対する計算グラフ生成機能(double backprop)をサポートしているため、同様の手順で高階微分を行うことができます。
ただし、vim-autogradにおける高階微分は、以下の二つの点に注意しなければなりません。

  1. 勾配の蓄積
    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の1階微分は、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.grad

    autograd#grad()

    let y = f(x)
let gx1 = autograd#grad(y, x, 1)
let gx2 = autograd#grad(gx1, x, 1)

    以上のように、その後にさらなる微分が控えている場合には、autograd#grad()関数で微分するか、backward()メソッドの直前にcleargrad()を呼びます。

  2. バックプロパゲーションのモード
    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)

以上に注意して、関数f(x)=x^5-2x^3+4x^2+6x+5の3階微分を求めてみます。

  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]

全体のコード

https://github.com/pit-ray/vim-autograd/blob/main/examples/higher-order.vim

ディープラーニングを利用したワイン分類

自動微分を用いると、損失関数の値が最小となるように、パラメータを最適化できます(勾配降下法)。ここではUCIで提供されているアルコールや酸の量などの13次元のベクトルから3種類のワインを分類するToyデータセットを利用して、ディープラーニングを行っていきます。データセットの詳細は以下の通りです。

出典: UCI Machine Learning Repository: Wine Data Set

https://archive.ics.uci.edu/ml/datasets/Wine

データセットの前処理

まず初めに、データセットを次のような形式の辞書として返す関数を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 indexes = autograd#shuffle(range(train_data_num))
    let epoch_loss = 0
    for l:i in range(each_iteration)
      let x = []
      let t = []
      for index in indexes[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ライクな記述でディープラーニングを行うことができました。

全体のコード

https://github.com/pit-ray/vim-autograd/blob/main/examples/wine-classify.vim

まとめと今後

Vim scriptで自動微分ライブラリを書いてみました。このライブラリを用いることで、高階微分やディープラーニングを行うことができます。今後、Vim9 scriptに対応することで、もう少し実用的なものになり、小さなニューラルネットワークを利用したプラグインなどを作成できるようになるかもしれません。

参考

本ライブラリのコア設計は、『ゼロから作る Deep Learning ❸』(O'Reilly Japan, 2020)のDeZeroをベースとしました。そのほかにも、逆伝播の実装にChainer、インターフェースの設計にPyTorchやNumPyを参考にしました。

また、ニューラルネットワークの学習をVim scriptで行う点に関連して、mattnさんのvim-brainを参考にしています。

https://github.com/mattn/vim-brain

脚注

  1. f(x)=xの微分は、f'(x)=1となります。同じ変数を2回使うg(x)=x+xの微分は、g'(x)=f'(x)+f'(x)となり、g'(x)=2と求められます。これは、加算の逆伝播が1であることからも、変数の使いまわしの結果として、勾配が加算されているといえます。 ↩︎

  2. vim-autogradが提供するFunctionを呼び出して合成関数の実装を行うより、逆伝播を個別に実装する方が効率が良いことが知られています。vim-autogradは、そのような利用にも対応しており、:help autograd#Function()に詳細があります。 ↩︎

  3. 当初、単純なSGDではうまく収束せず、あれやこれやと機能を付け足した結果、複雑なものができてしまいました。実装を見直したら、単純なSGDでも収束しました... ↩︎

AtCoder Beginner Contest 219
【macOS Monterey】HomebrewでPHPの実行環境をセットアップする