Class: GRX::Optim::Adam

Inherits:
Object
  • Object
show all
Defined in:
lib/grx/optim.rb

Overview

================================================================

Adam — Adaptive Moment Estimation (Kingma & Ba, 2015) The standard optimizer for deep neural networks.

Instance Method Summary collapse

Constructor Details

#initialize(params, lr: 0.001, betas: nil, beta1: 0.9, beta2: 0.999, eps: nil, epsilon: 1e-8, weight_decay: 0.0) ⇒ Adam

Returns a new instance of Adam.



57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
 
# File 'lib/grx/optim.rb', line 57

def initialize(params, lr: 0.001, betas: nil, beta1: 0.9, beta2: 0.999,
               eps: nil, epsilon: 1e-8, weight_decay: 0.0)
  @params       = params
  @lr           = lr.to_f
  if betas
    @beta1 = betas[0].to_f
    @beta2 = betas[1].to_f
  else
    @beta1 = beta1.to_f
    @beta2 = beta2.to_f
  end
  @epsilon      = eps ? eps.to_f : epsilon.to_f
  @weight_decay = weight_decay.to_f
  @t            = 0  # current step

  # First and second order moment vectors (zero-initialized)
  @m = params.map { |p| Tensor.zeros_like(p) }
  @v = params.map { |p| Tensor.zeros_like(p) }
end

Instance Method Details

#step ⇒ Object 



77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
 
# File 'lib/grx/optim.rb', line 77

def step
  @t += 1
  beta1t = @beta1 ** @t  # beta1^t for bias correction
  beta2t = @beta2 ** @t

  @params.each_with_index do |param, i|
    next unless param.grad

    grad = param.grad

    if @weight_decay > 0
      grad = grad + param.scale(@weight_decay)
    end

    if CAPI::LOADED
      CAPI.grx_adam_step(
        param.storage.ptr,
        @m[i].storage.ptr,
        @v[i].storage.ptr,
        grad.storage.ptr,
        @lr, @beta1, @beta2, @epsilon,
        beta1t, beta2t,
        param.numel
      )
    else
      # Pure Ruby fallback
      p_data = param.to_a
      m_data = @m[i].to_a
      v_data = @v[i].to_a
      g_data = grad.to_a
      p_data.each_with_index do |_, j|
        m_data[j] = @beta1 * m_data[j] + (1 - @beta1) * g_data[j]
        v_data[j] = @beta2 * v_data[j] + (1 - @beta2) * g_data[j] ** 2
        mh = m_data[j] / (1 - beta1t)
        vh = v_data[j] / (1 - beta2t)
        p_data[j] -= @lr * mh / (Math.sqrt(vh) + @epsilon)
      end
      param.storage.instance_variable_set(:@data, p_data)
      @m[i].storage.instance_variable_set(:@data, m_data)
      @v[i].storage.instance_variable_set(:@data, v_data)
    end
  end
end

#zero_grad ⇒ Object 



121
122
123
 
# File 'lib/grx/optim.rb', line 121

def zero_grad
  @params.each(&:zero_grad!)
end