Class: Torch::Optim::SGD

Inherits:
Optimizer show all
Defined in:
lib/torch/optim/sgd.rb

Instance Method Summary collapse

Constructor Details

#initialize(params, lr:) ⇒ SGD



4
5
6
7
# File 'lib/torch/optim/sgd.rb', line 4

def initialize(params, lr:)
  @params = params
  @lr = lr
end

Instance Method Details

#stepObject



18
19
20
21
22
23
24
25
# File 'lib/torch/optim/sgd.rb', line 18

def step
  @params.each do |param|
    next unless param.grad
    d_p = param.grad.data
    # same as param.data.add!(-@lr, d_p)
    param.data.sub!(d_p * @lr)
  end
end

#zero_gradObject



9
10
11
12
13
14
15
16
# File 'lib/torch/optim/sgd.rb', line 9

def zero_grad
  @params.each do |param|
    if param.grad
      param.grad.detach!
      param.grad.zero!
    end
  end
end