Is it necessary to make initialization for torch.nn modules


(Derek Lee) #1

Is it necessary to make initialization for torch.nn modules, such as GRU, Linear, LSTM, etc?

How to initialize the parameters if it is an important step? A complete example?


#2

You could use a custom initialization, if you don’t want to use the default ones.
The default initialization schemes are defined in the reset_parameters method of the module.
You can find such a method for nn.Conv2d here.

Here is an example on how to write your own initialization for a small CNN using xavier_normal:

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 6 ,3, 1, 1)
        self.pool1 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(6*12*12, 10)
        
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        return x

def weight_init(m):
    if isinstance(m, nn.Conv2d):
        nn.init.xavier_normal_(m.weight, gain=nn.init.calculate_gain('relu'))
        nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Linear):
        nn.init.xavier_normal_(m.weight)
        nn.init.zeros_(m.bias)

model = MyModel()
model.apply(weight_init)