Class: Daimond::NN::Conv2dRust

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

Instance Attribute Summary collapse

Instance Method Summary collapse

Methods inherited from Module

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

Constructor Details

#initialize(in_channels, out_channels, kernel_size) ⇒ Conv2dRust

Returns a new instance of Conv2dRust.



8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# File 'lib/daimond/nn/conv2d_rust.rb', line 8

def initialize(in_channels, out_channels, kernel_size)
  super()
  @in_channels = in_channels
  @out_channels = out_channels
  @kernel_size = kernel_size

  # Xavier инициализация
  k = kernel_size
  limit = Math.sqrt(2.0 / (in_channels * k * k))

  @weight = Tensor.new(
    Numo::DFloat.new(out_channels, in_channels, k, k).rand * 2 * limit - limit
  )
  @bias = Tensor.zeros(out_channels)
  @parameters = [@weight, @bias]
end

Instance Attribute Details

#biasObject (readonly)

Returns the value of attribute bias.



6
7
8
# File 'lib/daimond/nn/conv2d_rust.rb', line 6

def bias
  @bias
end

#weightObject (readonly)

Returns the value of attribute weight.



6
7
8
# File 'lib/daimond/nn/conv2d_rust.rb', line 6

def weight
  @weight
end

Instance Method Details

#forward(input) ⇒ Object



25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
# File 'lib/daimond/nn/conv2d_rust.rb', line 25

def forward(input)
  # input: [batch, in_c, h, w]
  batch = input.shape[0]
  in_c = @in_channels
  out_c = @out_channels
  h = input.shape[2]
  w = input.shape[3]
  k = @kernel_size

  # Используем Rust backend
  if Daimond::RustBackend.available?
    output_data = Daimond::RustBackend.conv2d(
      input.data, @weight.data, @bias.data,
      batch, in_c, out_c, h, w, k
    )
    out = Tensor.new(output_data, prev: [input, @weight, @bias], op: 'conv2d_rust')

    # Backward будет позже, пока заглушка
    out._backward = lambda {}

    return out
  else
    raise "Rust backend required for Conv2dRust"
  end
end