Class: Gsplat::Optim::Adam

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

Overview

Dense Adam optimizer with editable first-axis state for densification.

Direct Known Subclasses

SelectiveAdam

Defined Under Namespace

Classes: Group, State

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(groups = nil, lr: 1e-3, betas: [0.9, 0.999], eps: 1e-15, **named_groups) ⇒ Adam

rubocop:disable Naming/MethodParameterName

Raises:

  • (ArgumentError) 


16
17
18
19
20
21
22
23
24
25
 
# File 'lib/gsplat/optim/adam.rb', line 16

def initialize(groups = nil, lr: 1e-3, betas: [0.9, 0.999], eps: 1e-15, **named_groups)
  # rubocop:enable Naming/MethodParameterName
  raise ArgumentError, "provide positional or named groups, not both" if groups && !named_groups.empty?

  groups ||= named_groups
  @beta1, @beta2 = betas
  validate_hyperparameters!(lr, eps)
  @groups = normalize_groups(groups, lr, eps)
  @states = {}
end

Instance Attribute Details

#beta1 ⇒ Object (readonly)

Returns the value of attribute beta1.



13
14
15
 
# File 'lib/gsplat/optim/adam.rb', line 13

def beta1
  @beta1
end

#beta2 ⇒ Object (readonly)

Returns the value of attribute beta2.



13
14
15
 
# File 'lib/gsplat/optim/adam.rb', line 13

def beta2
  @beta2
end

#groups ⇒ Object (readonly)

Returns the value of attribute groups.



13
14
15
 
# File 'lib/gsplat/optim/adam.rb', line 13

def groups
  @groups
end

Instance Method Details

#append!(count) ⇒ Object

Raises:

  • (ArgumentError) 


108
109
110
111
112
113
114
115
116
117
118
119
120
121
 
# File 'lib/gsplat/optim/adam.rb', line 108

def append!(count)
  raise ArgumentError, "append count must be non-negative" unless count.is_a?(Integer) && !count.negative?
  return self if count.zero?

  groups.each_key do |name|
    current = state_for(name)
    @states[name] = State.new(
      step: current.step,
      exp_avg: append_zeros(current.exp_avg, count),
      exp_avg_sq: append_zeros(current.exp_avg_sq, count)
    )
  end
  self
end

#learning_rate(name = nil) ⇒ Numeric, Hash{Symbol=>Numeric}

Returns the current learning rate for one or all groups.

Returns:

  • (Numeric, Hash{Symbol=>Numeric}) 


69
70
71
72
73
74
 
# File 'lib/gsplat/optim/adam.rb', line 69

def learning_rate(name = nil)
  return groups.fetch(name.to_sym).lr if name
  return groups.values.first.lr if groups.length == 1

  groups.transform_values(&:lr)
end

#learning_rate=(value) ⇒ Object

Raises:

  • (ArgumentError) 


76
77
78
79
80
81
82
 
# File 'lib/gsplat/optim/adam.rb', line 76

def learning_rate=(value)
  raise ArgumentError, "learning rate must be positive" unless value.positive?

  @groups = groups.transform_values do |group|
    Group.new(name: group.name, variable: group.variable, lr: value, eps: group.eps)
  end
end

#load_state!(name = nil, step:, exp_avg:, exp_avg_sq:) ⇒ Object

Restores one parameter group's moments from a checkpoint.

Raises:

  • (ArgumentError) 


138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
 
# File 'lib/gsplat/optim/adam.rb', line 138

def load_state!(name = nil, step:, exp_avg:, exp_avg_sq:)
  key = name ? name.to_sym : single_group_name
  group = groups.fetch(key)
  expected = group.variable.data.shape
  unless exp_avg.shape == expected && exp_avg_sq.shape == expected
    raise ShapeError,
          "Adam state for #{key} expected #{expected.inspect}, " \
          "got #{exp_avg.shape.inspect} and #{exp_avg_sq.shape.inspect}"
  end
  raise ArgumentError, "Adam step must be a non-negative integer" unless step.is_a?(Integer) && !step.negative?

  type = group.variable.data.class
  @states[key] = State.new(
    step: step,
    exp_avg: type.cast(exp_avg).dup,
    exp_avg_sq: type.cast(exp_avg_sq).dup
  )
  self
end

#select!(selection) ⇒ self

Selects first-axis optimizer rows after pruning parameters.

Parameters:

  • selection (Numo::Bit, Array<Integer>, Range)

Returns:

  • (self) 


96
97
98
99
100
101
102
103
104
105
106
 
# File 'lib/gsplat/optim/adam.rb', line 96

def select!(selection)
  groups.each_key do |name|
    current = state_for(name)
    @states[name] = State.new(
      step: current.step,
      exp_avg: select_first_axis(current.exp_avg, selection),
      exp_avg_sq: select_first_axis(current.exp_avg_sq, selection)
    )
  end
  self
end

#set_learning_rate(name, value) ⇒ Object

Raises:

  • (ArgumentError) 


84
85
86
87
88
89
90
 
# File 'lib/gsplat/optim/adam.rb', line 84

def set_learning_rate(name, value)
  raise ArgumentError, "learning rate must be positive" unless value.positive?

  key = name.to_sym
  group = groups.fetch(key)
  @groups[key] = Group.new(name: key, variable: group.variable, lr: value, eps: group.eps)
end

#state(name = nil) ⇒ State, Hash{Symbol=>State}

Returns lazily initialized moment state for one or all groups.

Returns:



59
60
61
62
63
64
 
# File 'lib/gsplat/optim/adam.rb', line 59

def state(name = nil)
  return state_for(name.to_sym) if name
  return state_for(groups.keys.first) if groups.length == 1

  groups.keys.to_h { |key| [key, state_for(key)] }
end

#step ⇒ Object

rubocop:disable Metrics/AbcSize



28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
 
# File 'lib/gsplat/optim/adam.rb', line 28

def step
  groups.each_value do |group|
    gradient = group.variable.grad
    next unless gradient

    state = state_for(group.name)
    next_step = state.step + 1
    exp_avg = (beta1 * state.exp_avg) + ((1 - beta1) * gradient)
    exp_avg_sq = (beta2 * state.exp_avg_sq) + ((1 - beta2) * (gradient**2))
    corrected_mean = exp_avg / (1 - (beta1**next_step))
    corrected_variance = exp_avg_sq / (1 - (beta2**next_step))
    update = group.lr * corrected_mean / ((corrected_variance**0.5) + group.eps)
    group.variable.data[] = group.variable.data - update
    @states[group.name] = State.new(
      step: next_step,
      exp_avg: exp_avg,
      exp_avg_sq: exp_avg_sq
    )
  end
  self
end

#zero_grad! ⇒ Object

rubocop:enable Metrics/AbcSize



51
52
53
54
 
# File 'lib/gsplat/optim/adam.rb', line 51

def zero_grad!
  groups.each_value { |group| group.variable.zero_grad! }
  self
end

#zero_state_at!(indices) ⇒ self

Clears moments at selected first-axis rows.

Parameters:

  • indices (Integer, Array<Integer>, Range)

Returns:

  • (self) 


127
128
129
130
131
132
133
134
135
 
# File 'lib/gsplat/optim/adam.rb', line 127

def zero_state_at!(indices)
  groups.each_key do |name|
    current = state_for(name)
    index = [indices] + Array.new(current.exp_avg.ndim - 1, true)
    current.exp_avg[*index] = 0
    current.exp_avg_sq[*index] = 0
  end
  self
end