Class: Daimond::NN::Flatten

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

Instance Method Summary collapse

Methods inherited from Module

#call, #load, #parameters, #save, #zero_grad

Constructor Details

#initialize(start_dim: 1, end_dim: -1)) ⇒ Flatten

Returns a new instance of Flatten.



6
7
8
9
10
11
 
# File 'lib/daimond/nn/flatten.rb', line 6

def initialize(start_dim: 1, end_dim: -1)
  super()
  @start_dim = start_dim
  @end_dim = end_dim
  @input_shape = nil
end

Instance Method Details

#forward(input) ⇒ Object 



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

def forward(input)
  @input_shape = input.shape.dup
  batch = input.shape[0]
  rest = input.shape[1..-1].inject(:*)

  out_data = input.data.reshape(batch, rest)
  out = Tensor.new(out_data, prev: [input], op: 'flatten')

  out._backward = lambda do
    input.grad += out.grad.reshape(*@input_shape)
  end

  out
end