Class: GRX::NN::Module

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

Overview

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

Module — Base class for all neural network layers

Instance Method Summary collapse

Instance Method Details

#call(*args) ⇒ Object

Subclasses implement forward computation



41
42
43
 
# File 'lib/grx/nn.rb', line 41

def call(*args)
  forward(*args)
end

#load_weights(path) ⇒ Object 



36
37
38
 
# File 'lib/grx/nn.rb', line 36

def load_weights(path)
  GRX::Serialization.load(self, path)
end

#parameters ⇒ Object

Returns all trainable parameters for optimizer registration



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

def parameters
  instance_variables.flat_map do |var|
    val = instance_variable_get(var)
    case val
    when Tensor    then val.requires_grad ? [val] : []
    when Module    then val.parameters
    when Array     then val.flat_map { |v|
      case v
      when Tensor then v.requires_grad ? [v] : []
      when Module then v.parameters
      else []
      end
    }
    else []
    end
  end
end

#save_weights(path) ⇒ Object 



32
33
34
 
# File 'lib/grx/nn.rb', line 32

def save_weights(path)
  GRX::Serialization.save(self, path)
end

#zero_grad ⇒ Object 



28
29
30
 
# File 'lib/grx/nn.rb', line 28

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