Add weight normalization implementation

[colesbury]

This adds forward "pre-hooks" which get called before the module's
forward() method. Weight norm is implemented as a hook which calculates
the weight variable from the weight_g and weight_v every iteration.

Based on @rtqichen implementation.

[colesbury]

cc @rtqichen

I went with an implementation based on hooks, which is similar to your original function decorator implementation, but can be pickled and works with DataParallel.

[hanzhanggit]

@rtqichen @colesbury I did not find the data dependent initialization(which is used in the original paper) for the weight norm. I thought it was important to the weight norm. Do we plan to have it implemented or any examples about it?

[xwuaustin]

@rtqichen Thanks for contribution for this code.
@hanzhanggit @soumith @colesbury @ALL
Hi everyone who read this post. I have some questions regarding to weight_norm. It would be great if you can help.
I tried to implement the weight_norm for each convolution and linear layer (check the code here https://github.com/xwuaustin/weight_norm/blob/master/cifar10_tutorial_weightNorm.py ). However, the training loss in CIFAR-10 seems no difference to the original setting (see the picture below) at the first 2 epochs.

Now questions:

1. Is there something wrong with the code I modified? I used the code from cifar10_tutorial in pyTorch. All I did is to add the weight norm at each layer.

import torch.nn.utils.weight_norm as weightNorm
class Net(nn.Module):
def init(self):
super(Net, self).init()
### we use weight normalization after each convolutions and linear transfrom
self.conv1 = weightNorm(nn.Conv2d(3, 6, 5),name = "weight")
#print (self.conv1._parameters.keys())
self.pool = nn.MaxPool2d(2, 2)
self.conv2 =weightNorm(nn.Conv2d(6, 16, 5),name = "weight")
self.fc1 = weightNorm(nn.Linear(16 * 5 * 5, 120),name = "weight")
self.fc2 = weightNorm(nn.Linear(120, 84),name = "weight")
self.fc3 = weightNorm(nn.Linear(84, 10),name = "weight")

2 Is the update of the weights and bias, namely 'weight_g', 'weight_v', using the formulation:

3. Can we do the initialization as the paper suggested?

Thanks. Looking for your response. :)

Answer: I think I can answer this myself.
The first question: the way to use weight_norm should be fine. I use minibatch 128, then WN works better.
The second question: the update does follow what the paper suggested
The third question: we could do the initialization by defining a function say weight_init and then do net.apply(weight_init). There are references for weight initialization.

[sajadn]

I wrote a function which does the data-dependent initialization for weight norm layers (now it's just for conv layers, but it's straight forward to adapt it).
https://gist.github.com/sajadn/f29b14cd52023b614c334708c2a4345f

[tfjgeorge]

@colesbury Hi! Could you kindly explain what is the advantage (if any) of implementing weight norm with hooks as you do, versus implementing it as a new layer type (nn.Module) ?

[colesbury]

@tfjgeorge An nn.Module probably would made sense, but I don't remember the details so I'm not sure. (It's been over four years)