Sha256: 8dcba81a2b43a1a6fad4db7f898fe5202423a74b6e613b3fd6f725f006a813ce
Contents?: true
Size: 1.41 KB
Versions: 9
Compression:
Stored size: 1.41 KB
Contents
module Torch
module NN
class ConvNd < Module
def initialize(in_channels, out_channels, kernel_size, stride, padding, dilation, transposed, output_padding, groups, bias, padding_mode)
super()
raise ArgumentError, "in_channels must be divisible by groups" if in_channels % groups != 0
raise ArgumentError, "out_channels must be divisible by groups" if out_channels % groups != 0
@in_channels = in_channels
@out_channels = out_channels
@kernel_size = kernel_size
@stride = stride
@padding = padding
@dilation = dilation
@transposed = transposed
@output_padding = output_padding
@groups = groups
@padding_mode = padding_mode
if transposed
@weight = Parameter.new(Tensor.new(in_channels, out_channels / groups, *kernel_size))
else
@weight = Parameter.new(Tensor.new(out_channels, in_channels / groups, *kernel_size))
end
if bias
@bias = Parameter.new(Tensor.new(out_channels))
else
raise NotImplementedError
end
reset_parameters
end
def reset_parameters
Init.kaiming_uniform!(@weight, a: Math.sqrt(5))
if @bias
fan_in, _ = Init._calculate_fan_in_and_fan_out(@weight)
bound = 1 / Math.sqrt(fan_in)
Init.uniform!(@bias, a: -bound, b: bound)
end
end
end
end
end
Version data entries
9 entries across 9 versions & 1 rubygems