Class: GRX::NN::LayerNorm

Inherits:
Module
  • Object
show all
Defined in:
lib/grx/nn.rb

Overview

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

LayerNorm — Layer normalization across channel dimensions

Instance Attribute Summary collapse

Instance Method Summary collapse

Methods inherited from Module

#call, #load_weights, #parameters, #save_weights, #zero_grad

Constructor Details

#initialize(normalized_shape, epsilon: 1e-5) ⇒ LayerNorm

Returns a new instance of LayerNorm.



245
246
247
248
249
250
251
252
 
# File 'lib/grx/nn.rb', line 245

def initialize(normalized_shape, epsilon: 1e-5)
  @normalized_shape = normalized_shape.is_a?(Array) ? normalized_shape : [normalized_shape]
  @dim              = @normalized_shape.reduce(1, :*)
  @epsilon          = epsilon

  @gamma = Tensor.ones(@normalized_shape, requires_grad: true)
  @beta  = Tensor.zeros(@normalized_shape, requires_grad: true)
end

Instance Attribute Details

#beta ⇒ Object (readonly)

Returns the value of attribute beta.



243
244
245
 
# File 'lib/grx/nn.rb', line 243

def beta
  @beta
end

#epsilon ⇒ Object (readonly)

Returns the value of attribute epsilon.



243
244
245
 
# File 'lib/grx/nn.rb', line 243

def epsilon
  @epsilon
end

#gamma ⇒ Object (readonly)

Returns the value of attribute gamma.



243
244
245
 
# File 'lib/grx/nn.rb', line 243

def gamma
  @gamma
end

#normalized_shape ⇒ Object (readonly)

Returns the value of attribute normalized_shape.



243
244
245
 
# File 'lib/grx/nn.rb', line 243

def normalized_shape
  @normalized_shape
end

Instance Method Details

#forward(x) ⇒ Object 



254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
 
# File 'lib/grx/nn.rb', line 254

def forward(x)
  batch_size = x.shape[0]
  x_data = x.to_a

  means = Array.new(batch_size) do |b|
    x_data.slice(b * @dim, @dim).sum / @dim.to_f
  end
  vars = Array.new(batch_size) do |b|
    m = means[b]
    x_data.slice(b * @dim, @dim).sum { |v| (v - m)**2 } / @dim.to_f
  end

  gamma_data = @gamma.to_a
  beta_data  = @beta.to_a
  norm_data  = Array.new(batch_size * @dim)

  batch_size.times do |b|
    m = means[b]
    inv_std = 1.0 / Math.sqrt(vars[b] + @epsilon)
    @dim.times do |j|
      norm_data[b * @dim + j] = gamma_data[j] * (x_data[b * @dim + j] - m) * inv_std + beta_data[j]
    end
  end

  out = Tensor.create(norm_data, x.shape)
  if x.requires_grad || @gamma.requires_grad || @beta.requires_grad
    out.requires_grad = true
    out._grafo_hijos.push(x, @gamma, @beta)
    g_param = @gamma; b_param = @beta; d = @dim; eps = @epsilon
    out.backward_fn = ->(g) {
      g_data = g.to_a
      grad_gamma = Array.new(d, 0.0)
      grad_beta  = Array.new(d, 0.0)
      grad_x     = Array.new(batch_size * d, 0.0)

      batch_size.times do |b|
        m = means[b]; v = vars[b]
        inv_std = 1.0 / Math.sqrt(v + eps)
        x_hat = Array.new(d) { |j| (x_data[b * d + j] - m) * inv_std }
        dl_dxhat = Array.new(d) { |j| g_data[b * d + j] * gamma_data[j] }
        sum_dl = dl_dxhat.sum
        sum_dl_x = dl_dxhat.zip(x_hat).sum { |a, c| a * c }

        d.times do |j|
          grad_gamma[j] += g_data[b * d + j] * x_hat[j]
          grad_beta[j]  += g_data[b * d + j]
          grad_x[b * d + j] = (inv_std / d.to_f) * (d.to_f * dl_dxhat[j] - sum_dl - x_hat[j] * sum_dl_x)
        end
      end

      g_param.agregar_gradiente(Tensor.create(grad_gamma, g_param.shape)) if g_param.requires_grad
      b_param.agregar_gradiente(Tensor.create(grad_beta, b_param.shape))   if b_param.requires_grad
      x.agregar_gradiente(Tensor.create(grad_x, x.shape))                 if x.requires_grad
    }
  end
  out
end

#to_s ⇒ Object 



312
313
314
 
# File 'lib/grx/nn.rb', line 312

def to_s
  "LayerNorm(#{@normalized_shape})"
end