Class: Ignis::AI::NN::LayerNorm

Inherits:
Module
  • Object
show all
Defined in:
lib/nnw/ai/nn/layer_norm.rb

Overview

Layer normalization: y = gamma * (x - mean) / sqrt(var + eps) + beta Normalizes along the last dimension(s).

Instance Attribute Summary collapse

Attributes inherited from Module

#training

Instance Method Summary collapse

Methods inherited from Module

#call, #eval!, #load_state_dict, #named_parameters, #num_parameters, #parameters, #state_dict, #to, #train!, #zero_grad!

Constructor Details

#initialize(normalized_shape, eps: 1e-5, device_id: 0) ⇒ LayerNorm

Returns a new instance of LayerNorm.

Parameters:

  • normalized_shape (Integer) —

    size of the last dimension

  • eps (Float) (defaults to: 1e-5) —

    epsilon for numerical stability

  • device_id (Integer) (defaults to: 0) 


20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
 
# File 'lib/nnw/ai/nn/layer_norm.rb', line 20

def initialize(normalized_shape, eps: 1e-5, device_id: 0)
  super()
  @normalized_shape = normalized_shape
  @eps = eps

  # Initialize weight (gamma) to ones
  weight_nv = Ignis::Shared::NvArray.new(shape: [normalized_shape],
                                        dtype: :float32, device_id: device_id)
  weight_nv.from_host(Array.new(normalized_shape, 1.0))
  @weight = register_parameter("weight",
             Tensor.new(data: weight_nv, requires_grad: true))

  # Initialize bias (beta) to zeros
  bias_nv = Ignis::Shared::NvArray.new(shape: [normalized_shape],
                                      dtype: :float32, device_id: device_id)
  bias_nv.from_host(Array.new(normalized_shape, 0.0))
  @bias = register_parameter("bias",
           Tensor.new(data: bias_nv, requires_grad: true))
end

Instance Attribute Details

#bias ⇒ Tensor (readonly)

Returns beta (shift).

Returns:

  • (Tensor) —

    beta (shift)



15
16
17
 
# File 'lib/nnw/ai/nn/layer_norm.rb', line 15

def bias
  @bias
end

#weight ⇒ Tensor (readonly)

Returns gamma (scale).

Returns:

  • (Tensor) —

    gamma (scale)



12
13
14
 
# File 'lib/nnw/ai/nn/layer_norm.rb', line 12

def weight
  @weight
end

Instance Method Details

#forward(x) ⇒ Tensor

Forward pass: applies layer normalization.

Parameters:

  • x (Tensor) —

    input tensor [*, normalized_shape]

Returns:

  • (Tensor) —

    normalized tensor



43
44
45
 
# File 'lib/nnw/ai/nn/layer_norm.rb', line 43

def forward(x)
  x.layer_norm(@weight, @bias, eps: @eps)
end

#to_s ⇒ String

Returns:

  • (String) 


48
49
50
 
# File 'lib/nnw/ai/nn/layer_norm.rb', line 48

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