Gradients vanishing despite using Kaiming initialization

user_15504988 0
Viewed 181

I was implementing a conv block in pytorch with activation function(prelu). I used Kaiming initilization to initialize all my weights and set all the bias to zero. However as I tested these blocks (by stacking 100 such conv and activation blocks on top of each other), I noticed that the output I am getting values of the order of 10^(-10). Is this normal, considering I am stacking upto 100 layers. Adding a small bias to each layer fixes the problem. But in Kaiming initialization the biases are supposed to be zero.

Here is the conv block code

from collections import Iterable

def convBlock(
    input_channels, output_channels, kernel_size=3, padding=None, activation="prelu"
):
    """
    Initializes a conv block using Kaiming Initialization
    """
    padding_par = 0
    if padding == "same":
        padding_par = same_padding(kernel_size)
    conv = nn.Conv2d(input_channels, output_channels, kernel_size, padding=padding_par)
    relu_negative_slope = 0.25
    act = None
    if activation == "prelu" or activation == "leaky_relu":
        nn.init.kaiming_normal_(conv.weight, a=relu_negative_slope, mode="fan_in")
        if activation == "prelu":
            act = nn.PReLU(init=relu_negative_slope)
        else:
            act = nn.LeakyReLU(negative_slope=relu_negative_slope)
    if activation == "relu":
        nn.init.kaiming_normal_(conv.weight, nonlinearity="relu")
        act = nn.ReLU()
    nn.init.constant_(conv.bias.data, 0)
    block = nn.Sequential(conv, act)
    return block


def flatten(lis):
    for item in lis:
        if isinstance(item, Iterable) and not isinstance(item, str):
            for x in flatten(item):
                yield x
        else:
            yield item


def Sequential(args):
    flattened_args = list(flatten(args))
    return nn.Sequential(*flattened_args)

This is the test Code

ls=[]
for i in range(100):
    ls.append(convBlock(3,3,3,"same"))
model=Sequential(ls)

test=np.ones((1,3,5,5))
model(torch.Tensor(test))

And the output I am getting is

tensor([[[[-1.7771e-10, -3.5088e-10,  5.9369e-09,  4.2668e-09,  9.8803e-10],
          [ 1.8657e-09, -4.0271e-10,  3.1189e-09,  1.5117e-09,  6.6546e-09],
          [ 2.4237e-09, -6.2249e-10, -5.7327e-10,  4.2867e-09,  6.0034e-09],
          [-1.8757e-10,  5.5446e-09,  1.7641e-09,  5.7018e-09,  6.4347e-09],
          [ 1.2352e-09, -3.4732e-10,  4.1553e-10, -1.2996e-09,  3.8971e-09]],

         [[ 2.6607e-09,  1.7756e-09, -1.0923e-09, -1.4272e-09, -1.1840e-09],
          [ 2.0668e-10, -1.8130e-09, -2.3864e-09, -1.7061e-09, -1.7147e-10],
          [-6.7161e-10, -1.3440e-09, -6.3196e-10, -8.7677e-10, -1.4851e-09],
          [ 3.1475e-09, -1.6574e-09, -3.4180e-09, -3.5224e-09, -2.6642e-09],
          [-1.9703e-09, -3.2277e-09, -2.4733e-09, -2.3707e-09, -8.7598e-10]],

         [[ 3.5573e-09,  7.8113e-09,  6.8232e-09,  1.2285e-09, -9.3973e-10],
          [ 6.6368e-09,  8.2877e-09,  9.2108e-10,  9.7531e-10,  7.0011e-10],
          [ 6.6954e-09,  9.1019e-09,  1.5128e-08,  3.3151e-09,  2.1899e-10],
          [ 1.2152e-08,  7.7002e-09,  1.6406e-08,  1.4948e-08, -6.0882e-10],
          [ 6.9930e-09,  7.3222e-09, -7.4308e-10,  5.2505e-09,  3.4365e-09]]]],
       grad_fn=<PreluBackward>)
pytorch gradient
1 Answers

Amazing question (and welcome to StackOverflow)! Research paper for quick reference.

TLDR

  • Try wider networks (64 channels)
  • Add Batch Normalization after activation (or even before, shouldn't make much difference)
  • Add residual connections (shouldn't improve much over batch norm, last resort)

Please check this out in this order and give a comment what (and if) any of that worked in your case (as I'm also curious).

Things you do differently

  • Your neural network is very deep, yet very narrow (81 parameters per layer only!)

Due to above, one cannot reliably create those weights from normal distribution as the sample is just too small.

Try wider networks, 64 channels or more

  • You are trying much deeper network than they did

Section: Comparison Experiments

We conducted comparisons on a deep but efficient model with 14 weight layers (actually 22 was also tested in comparison with Xavier)

That was due to date of release of this paper (2015) and hardware limitations "back in the days" (let's say)

Is this normal?

Approach itself is quite strange with layers of this depth, at least currently;

  • each conv block is usually followed by activation like ReLU and Batch Normalization (which normalizes signal and helps with exploding/vanishing signals)
  • usually networks of this depth (even of depth half of what you've got) use also residual connections (though this is not directly linked to vanishing/small signal, more connected to degradation problem of even deep networks, like 1000 layers)
user_108864200
Related