def __init__(self, num_features): super().__init__() self.num_features = num_features self.weight = jt.zeros((num_features)) self.bias = jt.zeros((num_features)) init.uniform_(self.weight, 0, 1) init.constant_(self.bias, 0.0)
def test_uniform(self): a = init.uniform(5, "float32") assert ((a > 0) & (a < 1)).all() a = init.uniform((2, 3), low=-1, high=1) assert ((a > -1) & (a < 1)).all() linear = nn.Linear(2, 2) init.uniform_(linear.weight) assert (linear.weight > 0).all() linear.weight.uniform_() assert (linear.weight > 0).all()
def uniform(size, var): if var is not None: bound = 1.0 / math.sqrt(size) init.uniform_(var, -bound, bound)
def glorot(var): if var is not None: stdv = math.sqrt(6.0 / (var.size(-2) + var.size(-1))) init.uniform_(var, -stdv, stdv)
def kaiming_uniform(var, fan, a): if var is not None: bound = math.sqrt(6 / ((1 + a**2) * fan)) init.uniform_(var, -bound, bound)
def XavierFill(tensor): """Caffe2 XavierFill Implementation""" size = reduce(operator.mul, tensor.shape, 1) fan_in = size / tensor.shape[0] scale = math.sqrt(3 / fan_in) return init.uniform_(tensor, -scale, scale)
def uniform_(x,l,r): x = jt.array([x]) init.uniform_(x,l,r) return x.numpy().item()