Module: GRX::Utils

Defined in:
lib/grx/utils.rb

Class Method Summary collapse

Class Method Details

.clip_grad_norm!(parameters, max_norm) ⇒ Object

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

clip_grad_norm! — Clips gradients of parameter collection so their combined L2 norm does not exceed max_norm. Prevents exploding gradient problems during deep network training.



10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
 
# File 'lib/grx/utils.rb', line 10

def self.clip_grad_norm!(parameters, max_norm)
  max_norm = max_norm.to_f
  total_norm_sq = 0.0
  parameters.each do |p|
    next unless p.grad
    total_norm_sq += p.grad.square.to_a.sum
  end
  total_norm = Math.sqrt(total_norm_sq)
  clip_coef = max_norm / (total_norm + 1e-6)
  if clip_coef < 1.0
    parameters.each do |p|
      next unless p.grad
      p.grad = p.grad.scale(clip_coef)
    end
  end
  total_norm
end