Class: Ignis::AI::NN::Module

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

Overview

Module — base class for all neural network layers.

Provides parameter management, state dict serialization, training/eval mode toggling, and automatic parameter collection.

Examples:

class MyModel < Ignis::AI::NN::Module
  def initialize
    super
    @linear = register_module("linear", Linear.new(768, 256))
  end

  def forward(x)
    @linear.call(x).relu
  end
end

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize ⇒ Module

Returns a new instance of Module.



26
27
28
29
30
 
# File 'lib/nnw/ai/nn/module.rb', line 26

def initialize
  @training = true
  @modules = {}    # name → Module
  @parameters_map = {} # name → Tensor
end

Instance Attribute Details

#training ⇒ Boolean (readonly)

Returns whether in training mode.

Returns:

  • (Boolean) —

    whether in training mode



24
25
26
 
# File 'lib/nnw/ai/nn/module.rb', line 24

def training
  @training
end

Instance Method Details

#call(*args, **kwargs, &block) ⇒ Tensor

Make modules callable. Forwards positional AND keyword args (e.g. the ‘mask:` passed to Block/MultiHeadAttention) through to #forward.

Returns:

  • (Tensor) 


41
42
43
 
# File 'lib/nnw/ai/nn/module.rb', line 41

def call(*args, **kwargs, &block)
  forward(*args, **kwargs, &block)
end

#eval! ⇒ self

Switch to eval mode (disables dropout, uses running stats for BN).

Returns:

  • (self) 


80
81
82
83
84
 
# File 'lib/nnw/ai/nn/module.rb', line 80

def eval!
  @training = false
  @modules.each_value { |m| m.eval! }
  self
end

#forward(*args) ⇒ Object

Override in subclass to define forward pass.

Raises:

  • (NotImplementedError) 


34
35
36
 
# File 'lib/nnw/ai/nn/module.rb', line 34

def forward(*args)
  raise NotImplementedError, "#{self.class}#forward must be implemented"
end

#load_state_dict(dict, strict: true) ⇒ self

Load state dict from Hash of → Ruby Array.

Parameters:

  • dict (Hash{String => Array})
  • strict (Boolean) (defaults to: true) —

    if true, raises on missing/unexpected keys

Returns:

  • (self) 


120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
 
# File 'lib/nnw/ai/nn/module.rb', line 120

def load_state_dict(dict, strict: true)
  params = named_parameters

  if strict
    missing = params.keys - dict.keys
    unexpected = dict.keys - params.keys
    raise KeyError, "Missing keys: #{missing}" unless missing.empty?
    raise KeyError, "Unexpected keys: #{unexpected}" unless unexpected.empty?
  end

  dict.each do |name, values|
    next unless params.key?(name)
    param = params[name]
    param.data.from_host(values)
  end

  self
end

#named_parameters(prefix: "") ⇒ Hash{String => Tensor}

Named parameters as a flat hash.

Parameters:

  • prefix (String) (defaults to: "") —

    prefix for nested module names

Returns:

  • (Hash{String => Tensor}) 


57
58
59
60
61
62
63
64
65
66
67
68
 
# File 'lib/nnw/ai/nn/module.rb', line 57

def named_parameters(prefix: "")
  result = {}
  @parameters_map.each do |name, param|
    key = prefix.empty? ? name : "#{prefix}.#{name}"
    result[key] = param
  end
  @modules.each do |name, mod|
    mod_prefix = prefix.empty? ? name : "#{prefix}.#{name}"
    result.merge!(mod.named_parameters(prefix: mod_prefix))
  end
  result
end

#num_parameters ⇒ Integer

Total parameter count.

Returns:

  • (Integer) 


141
142
143
 
# File 'lib/nnw/ai/nn/module.rb', line 141

def num_parameters
  parameters.sum(&:numel)
end

#parameters ⇒ Array<Tensor>

Collect all leaf parameters (requires_grad: true) recursively.

Returns:

  • (Array<Tensor>) 


47
48
49
50
51
52
 
# File 'lib/nnw/ai/nn/module.rb', line 47

def parameters
  params = []
  @parameters_map.each_value { |p| params << p }
  @modules.each_value { |m| params.concat(m.parameters) }
  params
end

#state_dict ⇒ Hash{String => Array}

Export state dict (parameter name → Ruby Array of values).

Returns:

  • (Hash{String => Array}) 


108
109
110
111
112
113
114
 
# File 'lib/nnw/ai/nn/module.rb', line 108

def state_dict
  result = {}
  named_parameters.each do |name, param|
    result[name] = param.to_host
  end
  result
end

#to(device_id:) ⇒ self

Move all parameters to a specific device.

Parameters:

  • device_id (Integer)

Returns:

  • (self) 


89
90
91
92
93
94
95
96
97
98
 
# File 'lib/nnw/ai/nn/module.rb', line 89

def to(device_id:)
  parameters.each do |p|
    next if p.device_id == device_id
    host_data = p.to_host
    new_nv = Ignis::Shared::NvArray.new(shape: p.shape, dtype: p.dtype, device_id: device_id)
    new_nv.from_host(host_data)
    p.instance_variable_set(:@data, new_nv)
  end
  self
end

#to_s ⇒ String

String representation.

Returns:

  • (String) 


147
148
149
150
151
152
153
154
 
# File 'lib/nnw/ai/nn/module.rb', line 147

def to_s
  parts = ["#{self.class.name}("]
  @modules.each do |name, mod|
    parts << "  (#{name}): #{mod}"
  end
  parts << ")"
  parts.join("\n")
end

#train! ⇒ self

Switch to training mode (enables dropout, batch norm updates).

Returns:

  • (self) 


72
73
74
75
76
 
# File 'lib/nnw/ai/nn/module.rb', line 72

def train!
  @training = true
  @modules.each_value { |m| m.train! }
  self
end

#zero_grad! ⇒ void

This method returns an undefined value.

Zero all gradients.



102
103
104
 
# File 'lib/nnw/ai/nn/module.rb', line 102

def zero_grad!
  parameters.each(&:zero_grad!)
end