Class: Ignis::AI::Optim::LRScheduler::LinearWarmup

Inherits:
Object
  • Object
show all
Defined in:
lib/nnw/ai/optim/lr_scheduler.rb

Overview

Linear warmup: linearly increase LR from 0 to base_lr over warmup_steps. After warmup, keep constant.

Instance Method Summary collapse

Constructor Details

#initialize(optimizer, warmup_steps:) ⇒ LinearWarmup

Returns a new instance of LinearWarmup.

Parameters:

  • optimizer (Base)
  • warmup_steps (Integer) 


14
15
16
17
18
 
# File 'lib/nnw/ai/optim/lr_scheduler.rb', line 14

def initialize(optimizer, warmup_steps:)
  @optimizer = optimizer
  @warmup_steps = warmup_steps
  @base_lr = optimizer.lr
end

Instance Method Details

#current_lr ⇒ Float

Returns:

  • (Float) 


36
37
38
 
# File 'lib/nnw/ai/optim/lr_scheduler.rb', line 36

def current_lr
  @optimizer.lr
end

#step ⇒ Float

Call after each optimizer step to update LR.

Returns:

  • (Float) —

    new learning rate



22
23
24
25
26
27
28
29
30
31
32
33
 
# File 'lib/nnw/ai/optim/lr_scheduler.rb', line 22

def step
  current = @optimizer.step_count
  if current < @warmup_steps
    factor = current.to_f / @warmup_steps
    new_lr = @base_lr * factor
    @optimizer.lr = new_lr
    new_lr
  else
    @optimizer.lr = @base_lr
    @base_lr
  end
end