CODEWITHSUNDEEP
python-3.xneural-networkpytorchtorch
lalaland
lalaland

Reputation: 439

Initialize weight in pytorch neural net

I have created this neural net:

class _netD(nn.Module):
    def __init__(self, num_classes=1, nc=1, ndf=64):
        super(_netD, self).__init__()
        self.num_classes = num_classes
        # nc is number of channels
        # num_classes is number of classes
        # ndf is the number of output channel at the first layer
        self.main = nn.Sequential(
            # input is (nc) x 28 x 28
            # conv2D(in_channels, out_channels, kernelsize, stride, padding)
            nn.Conv2d(nc, ndf , 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 14 x 14
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 7 x 7
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 3 x 3
            nn.Conv2d(ndf * 4, ndf * 8, 3, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*8) x 2 x 2
            nn.Conv2d(ndf * 8, num_classes, 2, 1, 0, bias=False),
            # out size = batch x num_classes x 1 x 1
        )

        if self.num_classes == 1:
            self.main.add_module('prob', nn.Sigmoid())
            # output = probability
        else:
            pass
            # output = scores

    def forward(self, input):

        output = self.main(input)
        return output.view(input.size(0), self.num_classes).squeeze(1)

I want to loop through the different layers and apply a weight initialization depending on the type of layer. I am trying to do the following:

D = _netD()

for name, param in D.named_parameters():
    if type(param) == nn.Conv2d:
      param.weight.normal_(...)

But that is not working. Can you please help me?

Thanks

Upvotes: 0

Views: 2480

Answers (1)

jhso
jhso

Reputation: 3283

type(param) will only return the actual datatype called a parameter for any type of weight or data in the model. Because named_parameters() doesn't return anything useful in the name either when used on an nn.sequential-based model, you need to look at the modules to see which layers are specifically related to the nn.Conv2d class using isinstance as such:

for layer in D.modules():
    if isinstance(layer, nn.Conv2d):
         layer.weight.data.normal_(...)

Or, the way that is recommended by Soumith Chintala himself, actually just loop through your main module itself:

for L,layer in D.main:
    if isisntance(layer,nn.Conv2d):
         layer.weight.data.normal_(..)

I actually prefer the first because you don't have to specify the exact nn.sequential module itself, and will search all possible modules in the model, but either one should do the job for you.

Upvotes: 1

Related Questions