各種パラメータ最適化手法の実装（SGD, Momentum, AdaGrad, Adam）

　今回は「ゼロから作るDeepLearning」で紹介されている各種パラメータ最適化手法を、書籍のPythonのサンプルコードをベースに、Rubyで実装してみました。

www.oreilly.co.jp

　各手法のロジックについては書籍で説明されていますので割愛します。また、前回の記事で書いたように、Rubyでは値の受け渡しが参照の値渡しになるので、パラメータのハッシュの各値は配列として保持する前提です。

SGD（確率的勾配降下法）

　SGDは前回の記事でもすでに使っていたのと同じで、別クラスとして分けただけのものです。

class SGD
  def initialize(lr: 0.01)
    @lr = lr
  end

  def update(params:, grads:)
    params.keys.each do |key|
      params[key][0] -= @lr * grads[key]
    end
  end
end

Momentum

　paramsのハッシュの各値を配列として扱っている以外は、PythonのコードをそのままRubyに置き換えています。インスタンス変数 v にはparamsと同じ構造の値を持ちますが、インスタンス内でしか使わないため、ハッシュの値は配列にはせずにそのままパラメータを保持しています。

　ゼロ行列の生成は Numo::NArray.zeros メソッドを使っています。

Numo::NArray.zeros
http://ruby-numo.github.io/narray/narray/Numo/NArray.html#zeros-class_method

class Momentum
  def initialize(lr: 0.01, momentum: 0.9)
    @lr = lr
    @momentum = momentum
    @v = nil
  end

  def update(params:, grads:)
    if @v.nil?
      @v = {}
      params.each do |key, value|
        @v[key] = Numo::DFloat.zeros(value.first.shape)
      end
    end

    params.keys.each do |key|
      @v[key] = @momentum * @v[key] - @lr * grads[key]
      params[key][0] += @v[key]
    end
  end
end

AdaGrad

　こちらもMomentumと同様にPythonのコードを置き換えています。

　行列に対しての平方根の計算は、Numo::DFloat::Math.sqrt メソッドを使っています。

Numo::DFloat::Math.sqrt
http://ruby-numo.github.io/narray/narray/Numo/DFloat/Math.html#sqrt-class_method

class AdaGrad
  def initialize(lr: 0.01)
    @lr = lr
    @h = nil
  end

  def update(params:, grads:)
    if @h.nil?
      @h = {}
      params.each do |key, value|
        @h[key] = Numo::DFloat.zeros(value.first.shape)
      end
    end

    params.keys.each do |key|
      @h[key] += grads[key] * grads[key]
      params[key][0] -= @lr * grads[key] / (Numo::DFloat::Math.sqrt(@h[key]) + 1e-7)
    end
  end
end

Adam

　こちらも同様の置き換えです。

class Adam
  def initialize(lr: 0.001, beta1: 0.9, beta2: 0.999)
    @lr = lr
    @beta1 = beta1
    @beta2 = beta2
    @iter = 0
    @m = nil
    @v = nil
  end

  def update(params:, grads:)
    if @m.nil?
      @m = {}
      @v = {}
      params.each do |key, value|
        @m[key] = Numo::DFloat.zeros(value.first.shape)
        @v[key] = Numo::DFloat.zeros(value.first.shape)
      end
    end

    @iter += 1
    lr_t = @lr * Numo::DFloat::Math.sqrt(1.0 - @beta2 ** @iter) / (1.0 - @beta1 ** @iter)

    params.keys.each do |key|
      @m[key] += (1 - @beta1) * (grads[key] - @m[key])
      @v[key] += (1 - @beta2) * (grads[key] ** 2 - @v[key])

      params[key][0] -= lr_t * @m[key] / (Numo::DFloat::Math.sqrt(@v[key]) + 1e-7)
    end
  end
end

MNISTデータセットによる最適化手法の比較

　上記の最適化手法の実装について、MNISTデータセットを用いた学習の比較を行います。こちらも基本的には書籍のPython実装をベースにしていて、5 層のニューラルネットワークで、各層 100 個のニューロンを持つ ネットワークという構成です。活性化関数にはReLUを用いています。

　まずは複数レイヤのネットワークの処理を行うためのMultiLayerNetクラスの実装です。以前のTwoLayerNetを、3層以上のネットワークにも対応させた形です。

require 'numo/narray'
require './layers.rb'

class MultiLayerNet
  def initialize(input_size:, hidden_size_list:, output_size:, activation: :relu, weight_init_std: :relu, weight_decay_lambda: 0)
    @input_size          = input_size
    @output_size         = output_size
    @hidden_size_list    = hidden_size_list
    @hidden_layer_num    = hidden_size_list.size
    @weight_decay_lambda = weight_decay_lambda
    @params              = {}

    # 重みの初期化
    init_weight(weight_init_std)

    # レイヤの生成
    activation_layer = {
      sigmoid: Sigmoid,
      relu:    Relu
    }
    @layers = {}
    (1..@hidden_layer_num).each do |idx|
      @layers["Affine#{idx}"] = Affine.new(w: @params["w#{idx}"], b: @params["b#{idx}"])
      @layers["Activation_function#{idx}"] = activation_layer[activation].new
    end

    idx = @hidden_layer_num + 1
    @layers["Affine#{idx}"] = Affine.new(w: @params["w#{idx}"], b: @params["b#{idx}"])

    @last_layer = SoftmaxWithLoss.new
  end

  def params
    @params
  end

  def init_weight(weight_init_std)
    all_size_list = [@input_size] + @hidden_size_list + [@output_size]
    (1..(all_size_list.size - 1)).each do |idx|
      scale = weight_init_std
      if %i(relu he).include?(weight_init_std)
        scale = Numo::DFloat::Math.sqrt(2.0 / all_size_list[idx - 1])
      elsif %i(sigmoid xavier).include?(weight_init_std)
        scale = Numo::DFloat::Math.sqrt(1.0 / all_size_list[idx - 1])
      end

      Numo::NArray.srand
      @params["w#{idx}"] = [scale * Numo::DFloat.new(all_size_list[idx - 1], all_size_list[idx]).rand_norm]
      @params["b#{idx}"] = [Numo::DFloat.zeros(all_size_list[idx])]
    end
  end

  def predict(x:)
    @layers.values.inject(x) do |x, layer|
      x = layer.forward(x: x)
    end
  end

  # x: 入力データ, t: 教師データ
  def loss(x:, t:)
    y = predict(x: x)

    weight_decay = 0
    (1..(@hidden_layer_num + 1)).each do |idx|
      w = @params["w#{idx}"].first
      weight_decay += 0.5 * @weight_decay_lambda * (w ** 2).sum
    end
    @last_layer.forward(x: y, t: t) + weight_decay
  end

  def accuracy(x:, t:)
    y = predict(x: x)
    y = y.max_index(1) % 10
    if t.ndim != 1
      t = t.max_index(1) % 10
    end

    y.eq(t).cast_to(Numo::UInt16).sum / x.shape[0].to_f
  end

  def gradient(x:, t:)
    # forward
    loss(x: x, t: t)

    # backward
    dout = 1
    dout = @last_layer.backward(dout: dout)

    layers = @layers.values.reverse
    layers.inject(dout) do |dout, layer|
      dout = layer.backward(dout: dout)
    end

    grads = {}
    (1..(@hidden_layer_num + 1)).each do |idx|
      grads["w#{idx}"] = @layers["Affine#{idx}"].dw + @weight_decay_lambda * @layers["Affine#{idx}"].w.first
      grads["b#{idx}"] = @layers["Affine#{idx}"].db
    end

    grads
  end
end

　そして各手法での学習と、結果のグラフ描画を行うためのスクリプトの実装です。基本的な処理は前回の誤差逆伝播法での学習処理と同じで、SGD以外にもMomentum、AdaGrad、Adamでの学習を行い、結果を比較しています。

require 'numo/gnuplot'
require './mnist.rb'
require './optimizers.rb'
require './multi_layer_net.rb'

# 0: MNISTデータの読み込み
x_train, t_train, x_test, t_test = load_mnist(normalize: true)

train_size = x_train.shape[0]
batch_size = 128
max_iterations = 1500

# 1: 実験の設定
optimizers = {
  sgd:      SGD.new,
  momentum: Momentum.new,
  adagrad:  AdaGrad.new,
  adam:     Adam.new
}

networks = {}
train_loss = {}
optimizers.each do |key, optimizer|
  networks[key]   = MultiLayerNet.new(input_size: 784, hidden_size_list: [100, 100, 100, 100], output_size: 10)
  train_loss[key] = []
end

# 2: 訓練の開始
max_iterations.times do |i|
  Numo::NArray.srand
  batch_mask = Numo::Int32.new(batch_size).rand(0, train_size)
  x_batch = x_train[batch_mask, true]
  t_batch = t_train[batch_mask]

  optimizers.each do |key, optimizer|
    grads = networks[key].gradient(x: x_batch, t: t_batch)
    optimizers[key].update(params: networks[key].params, grads: grads)

    loss = networks[key].loss(x: x_batch, t: t_batch)
    train_loss[key] << loss
  end

  next unless i % 100 == 0

  puts "========== iteration: #{i} =========="
  optimizers.keys.each do |key|
    loss = networks[key].loss(x: x_batch, t: t_batch)
    puts "#{key}: #{loss}"
  end
end

# 3: グラフの描画
x = (0..(max_iterations - 1)).to_a
Numo.gnuplot do
  set xlabel: 'iterations'
  set ylabel: 'loss'
  set yrange: 0...1
  plot x, train_loss[:sgd],      { w: :lines, t: 'SGD',      lc_rgb: 'green',  lw: 1 },
       x, train_loss[:momentum], { w: :lines, t: 'Momentum', lc_rgb: 'orange', lw: 1 },
       x, train_loss[:adagrad],  { w: :lines, t: 'AdaGrad',  lc_rgb: 'red',    lw: 1 },
       x, train_loss[:adam],     { w: :lines, t: 'Adam',     lc_rgb: 'blue',   lw: 1 }
end

　これをirbから実行すると下記のようにコンソールに結果が100イテレーションごとに表示され、最後にグラフが描画されます。

irb(main):001:0> load './optimizer_compare_mnist.rb'
========== iteration: 0 ==========
sgd: 2.490228492804417
momentum: 2.492025460726973
adagrad: 2.0825493921580134
adam: 2.276848398313694
========== iteration: 100 ==========
sgd: 1.5001470235084333
momentum: 0.42706984858496944
adagrad: 0.2004084345838115
adam: 0.34096154614776963
========== iteration: 200 ==========
sgd: 0.7664658855425125
momentum: 0.2733076953685949
adagrad: 0.11048856792711875
adam: 0.24828428524427817
========== iteration: 300 ==========
sgd: 0.5159466996794285
momentum: 0.25594092908625543
adagrad: 0.13587365073017388
adam: 0.20545236418926946
========== iteration: 400 ==========
sgd: 0.5042479159281286
momentum: 0.24538385847033825
adagrad: 0.11124732792005954
adam: 0.1797581753729203
========== iteration: 500 ==========
sgd: 0.32290967978019125
momentum: 0.1599522422679423
adagrad: 0.05731788233265379
adam: 0.0888823836035264
========== iteration: 600 ==========
sgd: 0.44467997494741673
momentum: 0.2578398459161452
adagrad: 0.1316129116477675
adam: 0.22439066383913017
========== iteration: 700 ==========
sgd: 0.28407622085704987
momentum: 0.10056311655844065
adagrad: 0.07989693533502204
adam: 0.0821317626167635
========== iteration: 800 ==========
sgd: 0.2706466278682429
momentum: 0.15550352523100197
adagrad: 0.06489312962717893
adam: 0.08528336870483003
========== iteration: 900 ==========
sgd: 0.2240422184352822
momentum: 0.11062658792202897
adagrad: 0.059913720263603615
adam: 0.03302573552710864
========== iteration: 1000 ==========
sgd: 0.3832020077768542
momentum: 0.13726781722583942
adagrad: 0.03417701415203686
adam: 0.053558781255080776
========== iteration: 1100 ==========
sgd: 0.38619949224379774
momentum: 0.15175966760909282
adagrad: 0.04222220798211423
adam: 0.06822940475295906
========== iteration: 1200 ==========
sgd: 0.2998755819916694
momentum: 0.07572742725924923
adagrad: 0.075962654676941
adam: 0.02748595912322749
========== iteration: 1300 ==========
sgd: 0.25322815781566416
momentum: 0.06003606774412698
adagrad: 0.03236788855975958
adam: 0.053987864918752786
========== iteration: 1400 ==========
sgd: 0.2720482764348912
momentum: 0.09105631835160209
adagrad: 0.044338756972504875
adam: 0.0668196066196452
=> true

　イテレーションの回数を1,500回としていますが、手元のVM環境ではこれ以上回数を増やすと途中で強制終了されてしまいました。Pythonコードをそのままの構成で移植しているので、もっとRubyに最適化した実装を考慮する必要がありそうです。

　Numo.gnuplot でのグラフ描画時は、setメソッドによる設定はplotより前に実行しておかないとグラフに反映されないようでした。

　コードは下記リポジトリにも公開しています。

github.com