Class: Ignis::AI::Optim::Base

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

Overview

Base optimizer class. Provides shared infrastructure for all optimizers.

Direct Known Subclasses

Adam, AdamW, SGD

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(params, lr:) ⇒ Base

Returns a new instance of Base.

Parameters:

  • params (Array<Tensor>) —

    parameters to optimize

  • lr (Float) —

    learning rate



17
18
19
20
21
 
# File 'lib/nnw/ai/optim/base.rb', line 17

def initialize(params, lr:)
  @params = params
  @lr = lr.to_f
  @step_count = 0
end

Instance Attribute Details

#params ⇒ Array<Tensor> (readonly)

Returns parameters being optimized.

Returns:

  • (Array<Tensor>) —

    parameters being optimized



10
11
12
 
# File 'lib/nnw/ai/optim/base.rb', line 10

def params
  @params
end

#step_count ⇒ Integer (readonly)

Returns current step count for LR scheduling.

Returns:

  • (Integer) —

    current step count for LR scheduling



13
14
15
 
# File 'lib/nnw/ai/optim/base.rb', line 13

def step_count
  @step_count
end

Instance Method Details

#clip_grad_norm!(max_norm) ⇒ Float

Clip gradient global norm across all parameters.

Parameters:

  • max_norm (Float)

Returns:

  • (Float) —

    the actual total norm (before clipping)



51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
 
# File 'lib/nnw/ai/optim/base.rb', line 51

def clip_grad_norm!(max_norm)
  # Phase 1: compute total squared sum across all params
  total_norm_nv = Ignis::Shared::NvArray.new(shape: [1], dtype: :float32,
                                            device_id: @params.first.device_id)
  total_norm_nv.from_host([0.0])

  @params.each do |p|
    next unless p.grad

    sq_kernel = Ignis::JIT::Kernels::Optimizer.grad_squared_sum
    n = p.numel
    block_size = 256
    grid_size = (n + block_size - 1) / block_size
    sq_kernel.launch(grid: [grid_size], block: [block_size],
                     shared_mem: block_size * 4,
                     args: [p.grad, total_norm_nv, n])
  end
  Ignis.synchronize

  total_norm = Math.sqrt(total_norm_nv.to_host[0])

  if total_norm > max_norm
    clip_factor = max_norm / (total_norm + 1e-6)
    clip_kernel = Ignis::JIT::Kernels::Optimizer.grad_clip_scale

    @params.each do |p|
      next unless p.grad
      n = p.numel
      clip_kernel.launch(grid: [(n + 255) / 256], block: [256],
                         args: [p.grad, clip_factor.to_f, n])
    end
    Ignis.synchronize
  end

  total_norm
end

#lr ⇒ Float

Get current learning rate.

Returns:

  • (Float) 


37
38
39
 
# File 'lib/nnw/ai/optim/base.rb', line 37

def lr
  @lr
end

#lr=(new_lr) ⇒ void

This method returns an undefined value.

Set learning rate (for schedulers).

Parameters:

  • new_lr (Float) 


44
45
46
 
# File 'lib/nnw/ai/optim/base.rb', line 44

def lr=(new_lr)
  @lr = new_lr.to_f
end

#step ⇒ void

This method returns an undefined value.

Perform one optimization step using accumulated gradients.

Raises:

  • (NotImplementedError) 


25
26
27
 
# File 'lib/nnw/ai/optim/base.rb', line 25

def step
  raise NotImplementedError, "Subclass must implement #step"
end

#zero_grad! ⇒ void

This method returns an undefined value.

Zero all parameter gradients.



31
32
33
 
# File 'lib/nnw/ai/optim/base.rb', line 31

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