Python torch.nn.init.calculate_gain() Examples

def __init__(self, in_shape, normalize):
        super().__init__()
        bias = not normalize
        self._nb_output_channel = 3136
        self.conv1 = Conv2d(in_shape[0], 32, 8, stride=4, padding=0, bias=bias)
        self.conv2 = Conv2d(32, 64, 4, stride=2, padding=0, bias=bias)
        self.conv3 = Conv2d(64, 64, 3, stride=1, padding=0, bias=bias)

        if normalize:
            self.bn1 = BatchNorm2d(32)
            self.bn2 = BatchNorm2d(64)
            self.bn3 = BatchNorm2d(64)
        else:
            self.bn1 = Identity()
            self.bn2 = Identity()
            self.bn3 = Identity()

        relu_gain = init.calculate_gain("relu")
        self.conv1.weight.data.mul_(relu_gain)
        self.conv2.weight.data.mul_(relu_gain)
        self.conv3.weight.data.mul_(relu_gain) 

def __init__(self, in_shape, normalize):
        super().__init__()
        bias = not normalize
        self._nb_output_channel = 2592
        self.conv1 = Conv2d(in_shape[0], 16, 8, stride=4, padding=0, bias=bias)
        self.conv2 = Conv2d(16, 32, 4, stride=2, padding=0, bias=bias)

        if normalize:
            self.bn1 = BatchNorm2d(16)
            self.bn2 = BatchNorm2d(32)
        else:
            self.bn1 = Identity()
            self.bn2 = Identity()

        relu_gain = init.calculate_gain("relu")
        self.conv1.weight.data.mul_(relu_gain)
        self.conv2.weight.data.mul_(relu_gain) 

def __init__(self, in_shape, normalize):
        super().__init__()
        bias = not normalize
        self._nb_output_channel = 3200
        self.conv1 = Conv2d(in_shape[0], 32, 7, stride=2, padding=1, bias=bias)
        self.conv2 = Conv2d(32, 64, 3, stride=2, padding=1, bias=bias)
        self.conv3 = Conv2d(64, 64, 3, stride=2, padding=1, bias=bias)
        self.conv4 = Conv2d(64, 128, 3, stride=2, padding=1, bias=bias)

        if normalize:
            self.bn1 = BatchNorm2d(32)
            self.bn2 = BatchNorm2d(64)
            self.bn3 = BatchNorm2d(64)
            self.bn4 = BatchNorm2d(128)
        else:
            self.bn1 = Identity()
            self.bn2 = Identity()
            self.bn3 = Identity()
            self.bn4 = Identity()

        relu_gain = init.calculate_gain("relu")
        self.conv1.weight.data.mul_(relu_gain)
        self.conv2.weight.data.mul_(relu_gain)
        self.conv3.weight.data.mul_(relu_gain)
        self.conv4.weight.data.mul_(relu_gain) 

def reset_parameters(self):
        # init real weight
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.A)

        # init A
        gain = calculate_gain('leaky_relu', 0)
        std = gain / np.sqrt(fan_in)
        bound = np.sqrt(3.0) * std

        with torch.no_grad():
            # TODO: find more stable initial values
            self.A.uniform_(-bound * (1 / (np.pi ** 2)), bound * (1 / (np.pi ** 2)))
            #
            # B is initialized by pi
            # -pi and pi is too big, so it is powed by -1
            self.B.uniform_(-1 / np.pi, 1 / np.pi) 

def _weight_initializer(tensor):
        # init.xavier_uniform_(tensor.weight, gain=init.calculate_gain('relu'))
        init.kaiming_normal_(tensor.weight, nonlinearity='relu')
        init.constant_(tensor.bias, 0.0) 

def init_param(self):
        if self.upsample:
            init.xavier_uniform_(self.W_branch.data, gain=init.calculate_gain('relu'))

        stdv = 1. / math.sqrt(self.out_feature)
        self.bias.data.uniform_(-stdv, stdv) 

def _weight_initializer(tensor):
        init.xavier_uniform_(tensor.weight, gain=init.calculate_gain('relu'))
        init.constant_(tensor.bias, 0.0) 

def _initialize_weights(self):
        init.orthogonal(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal(self.conv4.weight) 

def __init__(self, c_in, c_out):
        super(EqualizedLinear, self).__init__()
        self.linear = nn.Linear(c_in, c_out, bias=False)
        kaiming_normal(self.linear.weight, a=calculate_gain('linear'))
        # Scaling the weights for equalized learning
        linear_w = self.linear.weight.data.clone()
        self.bias = torch.nn.Parameter(torch.FloatTensor(c_out).fill_(0))
        self.scale = (torch.mean(self.linear.weight.data ** 2)) ** 0.5
        self.linear.weight.data.copy_(self.linear.weight.data / self.scale) 

def __init__(self, c_in, c_out, k_size, stride, pad):
        super(EqualizedDeconv2d, self).__init__()
        self.deconv = nn.ConvTranspose2d(c_in, c_out, k_size, stride, pad, bias=False)
        kaiming_normal(self.deconv.weight, a=calculate_gain('conv2d'))
        # Scaling the weights for equalized learning
        deconv_w = self.deconv.weight.data.clone()
        self.bias = torch.nn.Parameter(torch.FloatTensor(c_out).fill_(0))
        self.scale = (torch.mean(self.deconv.weight.data ** 2)) ** 0.5
        self.deconv.weight.data.copy_(self.deconv.weight.data / self.scale) 

def __init__(self, c_in, c_out, k_size, stride, pad, bias=False):
        super(EqualizedConv2d, self).__init__()
        self.conv = nn.Conv2d(c_in, c_out, k_size, stride, pad, bias=False)
        kaiming_normal(self.conv.weight, a=calculate_gain('conv2d'))
        # Scaling the weights for equalized learning
        conv_w = self.conv.weight.data.clone()
        self.bias = torch.nn.Parameter(torch.FloatTensor(c_out).fill_(0))
        self.scale = (torch.mean(self.conv.weight.data ** 2)) ** 0.5
        self.conv.weight.data.copy_(self.conv.weight.data / self.scale)     # for equalized learning rate 

def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight) 

def reset_parameters(self) -> None:
        # Embeddings
        for name in 'word pos nt action'.split():
            embedding = getattr(self, f'{name}_embedding')
            embedding.reset_parameters()

        # Encoders
        for name in 'stack buffer history'.split():
            encoder = getattr(self, f'{name}_encoder')
            encoder.reset_parameters()

        # Compositions
        for name in 'fwd bwd'.split():
            lstm = getattr(self, f'{name}_composer')
            for pname, pval in lstm.named_parameters():
                if pname.startswith('weight'):
                    init.orthogonal(pval)
                else:
                    assert pname.startswith('bias')
                    init.constant(pval, 0.)

        # Transformations
        gain = init.calculate_gain('relu')
        for name in 'word nt action'.split():
            layer = getattr(self, f'{name}2encoder')
            init.xavier_uniform(layer[0].weight, gain=gain)
            init.constant(layer[0].bias, 1.)
        init.xavier_uniform(self.fwdbwd2composed[0].weight, gain=gain)
        init.constant(self.fwdbwd2composed[0].bias, 1.)
        init.xavier_uniform(self.encoders2summary[1].weight, gain=gain)
        init.constant(self.encoders2summary[1].bias, 1.)
        init.xavier_uniform(self.summary2actionlogprobs.weight)
        init.constant(self.summary2actionlogprobs.bias, 0.)

        # Guards
        for name in 'stack buffer history'.split():
            guard = getattr(self, f'{name}_guard')
            init.constant(guard, 0.) 

def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight) 

def __init__(self, c_in, c_out, initializer='kaiming'):
        super(equalized_linear, self).__init__()
        self.linear = nn.Linear(c_in, c_out, bias=False)
        if initializer == 'kaiming':    kaiming_normal(self.linear.weight, a=calculate_gain('linear'))
        elif initializer == 'xavier':   torch.nn.init.xavier_normal(self.linear.weight)
        
        linear_w = self.linear.weight.data.clone()
        self.bias = torch.nn.Parameter(torch.FloatTensor(c_out).fill_(0))
        self.scale = (torch.mean(self.linear.weight.data ** 2)) ** 0.5
        self.linear.weight.data.copy_(self.linear.weight.data/self.scale) 

def __init__(self, c_in, c_out, k_size, stride, pad, initializer='kaiming'):
        super(equalized_deconv2d, self).__init__()
        self.deconv = nn.ConvTranspose2d(c_in, c_out, k_size, stride, pad, bias=False)
        if initializer == 'kaiming':    kaiming_normal(self.deconv.weight, a=calculate_gain('conv2d'))
        elif initializer == 'xavier':   xavier_normal(self.deconv.weight)
        
        deconv_w = self.deconv.weight.data.clone()
        self.bias = torch.nn.Parameter(torch.FloatTensor(c_out).fill_(0))
        self.scale = (torch.mean(self.deconv.weight.data ** 2)) ** 0.5
        self.deconv.weight.data.copy_(self.deconv.weight.data/self.scale) 

def __init__(self, c_in, c_out, k_size, stride, pad, initializer='kaiming', bias=False):
        super(equalized_conv2d, self).__init__()
        self.conv = nn.Conv2d(c_in, c_out, k_size, stride, pad, bias=False)
        if initializer == 'kaiming':    kaiming_normal(self.conv.weight, a=calculate_gain('conv2d'))
        elif initializer == 'xavier':   xavier_normal(self.conv.weight)
        
        conv_w = self.conv.weight.data.clone()
        self.bias = torch.nn.Parameter(torch.FloatTensor(c_out).fill_(0))
        self.scale = (torch.mean(self.conv.weight.data ** 2)) ** 0.5
        self.conv.weight.data.copy_(self.conv.weight.data/self.scale) 

def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight) 

def __init__(self, in_shape, nb_head, normalize):
        self._nb_output_channel = 800
        super().__init__()
        self.normalize = normalize
        bias = not normalize
        self.conv1 = Conv2d(
            in_shape[0], 32, kernel_size=3, stride=2, padding=1, bias=bias
        )
        self.conv2 = Conv2d(
            32, 32, kernel_size=3, stride=2, padding=1, bias=bias
        )

        self.attention = MultiHeadSelfAttention(20 * 20, 34, 34, nb_head)
        self.mlp = Linear(34, 34)

        self.conv3 = Conv2d(
            34, 32, kernel_size=3, stride=2, padding=1, bias=bias
        )
        self.conv4 = Conv2d(
            32, 32, kernel_size=3, stride=2, padding=1, bias=bias
        )

        if normalize:
            self.bn1 = BatchNorm2d(32)
            self.bn2 = BatchNorm2d(32)
            self.bn3 = BatchNorm2d(32)
            self.bn4 = BatchNorm2d(32)
        else:
            self.bn1 = Identity()
            self.bn2 = Identity()
            self.bn3 = Identity()
            self.bn4 = Identity()

        relu_gain = init.calculate_gain("relu")
        self.conv1.weight.data.mul_(relu_gain)
        self.conv2.weight.data.mul_(relu_gain)
        self.conv3.weight.data.mul_(relu_gain)
        self.conv4.weight.data.mul_(relu_gain) 

def __init__(self, in_shape, id, normalize):
        super().__init__(in_shape, id)
        bias = not normalize
        self._in_shape = in_shape
        self._out_shape = None
        self.conv1 = Conv2d(in_shape[0], 32, 7, stride=2, padding=1, bias=bias)
        self.conv2 = Conv2d(32, 32, 3, stride=2, padding=1, bias=bias)
        self.conv3 = Conv2d(32, 32, 3, stride=2, padding=1, bias=bias)
        self.conv4 = Conv2d(32, 32, 3, stride=2, padding=1, bias=bias)

        if normalize == "bn":
            self.bn1 = BatchNorm2d(32)
            self.bn2 = BatchNorm2d(32)
            self.bn3 = BatchNorm2d(32)
            self.bn4 = BatchNorm2d(32)
        elif normalize == "gn":
            self.bn1 = GroupNorm(8, 32)
            self.bn2 = GroupNorm(8, 32)
            self.bn3 = GroupNorm(8, 32)
            self.bn4 = GroupNorm(8, 32)
        else:
            self.bn1 = Identity()
            self.bn2 = Identity()
            self.bn3 = Identity()
            self.bn4 = Identity()

        relu_gain = init.calculate_gain("relu")
        self.conv1.weight.data.mul_(relu_gain)
        self.conv2.weight.data.mul_(relu_gain)
        self.conv3.weight.data.mul_(relu_gain)
        self.conv4.weight.data.mul_(relu_gain) 

def _weight_init(m):
        if isinstance(m, nn.Conv2d):
            init.orthogonal_(m.weight, init.calculate_gain('relu')) 

def xavier_uniform_sigmoid(modules):
    for m in modules:
        if isinstance(m, nn.Conv2d):
            init.xavier_uniform(m.weight.data, gain=init.calculate_gain('sigmoid'))
            if m.bias is not None:
                m.bias.data.zero_()
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.fill_(1)
            m.bias.data.zero_() 

def xavier_uniform_relu(modules):
    for m in modules:
        if isinstance(m, nn.Conv2d):
            init.xavier_uniform(m.weight.data, gain=init.calculate_gain('relu'))
            if m.bias is not None:
                m.bias.data.zero_()
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.fill_(1)
            m.bias.data.zero_() 

def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight) 

def reset_parameters(self):
        """Reinitialize learnable parameters."""
        if self.bias is not None:
            init.zeros_(self.bias)
        for i in range(self._k):
            init.xavier_normal_(self.W[i], init.calculate_gain('relu')) 

def reset_parameters(self):
        """Reinitialize learnable parameters."""
        gain = init.calculate_gain('relu')
        self.gru.reset_parameters()
        for linear in self.linears:
            init.xavier_normal_(linear.weight, gain=gain)
            init.zeros_(linear.bias) 

def reset_parameters(self):
        """Reinitialize learnable parameters."""
        gain = init.calculate_gain('relu')
        if self.bias is not None:
            nn.init.zeros_(self.bias)
        if isinstance(self.res_fc, nn.Linear):
            nn.init.xavier_normal_(self.res_fc.weight, gain=gain) 

def reset_parameters(self):
        """Reinitialize learnable parameters."""
        gain = init.calculate_gain('relu')
        init.xavier_normal_(self.fc.weight, gain=gain)
        if isinstance(self.res_fc, nn.Linear):
            init.xavier_normal_(self.res_fc.weight, gain=gain)
        init.normal_(self.mu.data, 0, 0.1)
        init.constant_(self.inv_sigma.data, 1)
        if self.bias is not None:
            init.zeros_(self.bias.data) 

def test_helper(self):


        class SuperResolutionNet(nn.Module):
            def __init__(self, upscale_factor, inplace=False):
                super(SuperResolutionNet, self).__init__()

                self.relu = nn.ReLU(inplace=inplace)
                self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
                self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
                self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
                self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
                self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

                self._initialize_weights()

            def forward(self, x):
                x = self.relu(self.conv1(x))
                x = self.relu(self.conv2(x))
                x = self.relu(self.conv3(x))
                x = self.pixel_shuffle(self.conv4(x))
                return x

            def _initialize_weights(self):
                init.orthogonal(self.conv1.weight, init.calculate_gain('relu'))
                init.orthogonal(self.conv2.weight, init.calculate_gain('relu'))
                init.orthogonal(self.conv3.weight, init.calculate_gain('relu'))
                init.orthogonal(self.conv4.weight)

        torch_model = SuperResolutionNet(upscale_factor=3)

        fake_input = Variable(torch.randn(1, 1, 224, 224), requires_grad=True)

        # use ModelHelper to create a C2 net
        helper = ModelHelper(name="test_model")
        start = helper.Sigmoid(['the_input'])
        # Embed the ONNX-converted pytorch net inside it
        toutput, = PyTorchModule(helper, torch_model, (fake_input,), [start])
        output = helper.Sigmoid(toutput)



        workspace.RunNetOnce(helper.InitProto())
        workspace.FeedBlob('the_input',fake_input.data.numpy())
        #print([ k for k in workspace.blobs ])
        workspace.RunNetOnce(helper.Proto())
        c2_out = workspace.FetchBlob(str(output))

        torch_out = torch.sigmoid(torch_model(torch.sigmoid(fake_input)))

        np.testing.assert_almost_equal(torch_out.data.cpu().numpy(), c2_out, decimal=3)