def reset_parameters(self):
self.weight.set_data(
initializer(HeUniform(math.sqrt(5)), self.weight.shape))
if self.has_bias:
fan_in, _ = _calculate_fan_in_and_fan_out(self.weight.shape)
bound = 1 / math.sqrt(fan_in)
self.bias.set_data(initializer(Uniform(bound),
[self.out_channels]))
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: Union[int, Tuple[int, ...]],
stride: Union[int, Tuple[int, ...]] = 1,
padding: Union[str, int, Tuple[int, ...]] = 0,
dilation: Union[int, Tuple[int, ...]] = 1,
groups: int = 1,
has_bias: bool = True,
padding_mode: str = 'zeros',
):
fan_in, _ = _calculate_fan_in_and_fan_out(self.weight.shape)
bound = 1 / math.sqrt(fan_in)
weight_init = HeUniform(math.sqrt(5))
bias_init = Uniform(bound)
super().__init__(in_channels, out_channels, kernel_size, stride,
padding, dilation, groups, has_bias, padding_mode,
weight_init, bias_init)
self.conv = ops.Conv2D(
out_channels,
kernel_size,
)