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, )