Hello all,
I am basically trying to do mean-variance normalization in embedding space. I have a RNN that embeds a sequence, and then calculates the mean and standard deviation. These are used to normalize the input of another sub-network. This model seems to work (atleast my training and validation losses behave themselves). I now wanted to try to learn weights corresponding to the mean and variance. Sort of like batch-norm (in the most hand-wavy way possible)
Q. Ofcourse, pytorch is magical 
and my model seems to be training. But am I doing it correctly?
Q. Do I need to do something else with the 2 new weight matrices I have introduced, or will they automatically be added to the parameter list of the RNN class?
Q. What would be the best way to check if these weights are actually learning something? some sort of check for the gradients w.r.t to these parameters?
My code :
class RNN(nn.Module):
def __init__(self, input_size, hidden_size, num_layers):
super(RNN, self).__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.input_size = input_size
self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
#project down to feature dimension
self.proj = nn.Linear(hidden_size, 840)
#Define weights for the channel mean and variance
self.W_mean = nn.Parameter(torch.randn(840,840))
self.W_std = nn.Parameter(torch.randn(840,840))
def forward(self,seq,x):
# Set initial states
h0 = Variable(torch.zeros(self.num_layers, seq.size(0), self.hidden_size).cuda())
c0 = Variable(torch.zeros(self.num_layers, seq.size(0), self.hidden_size).cuda())
out, _ = self.lstm(seq, (h0,c0))
#project lstm embeddings to feature size
out = out.contiguous().view(out.size(0)*out.size(1),self.hidden_size)
proj = self.proj(out)
proj = proj.view(-1,seq.size(1),x.size(1))
#840 dimensional average embedding
avg_emb = torch.mean(proj,1)
std_emb = torch.std(proj,dim=1)
#subtract the avg embedding from the speech frames
avg_emb = avg_emb.view(-1,x.size(1))
std_emb = std_emb.view(-1,x.size(1))
#mean vaiance normalization
x_norm = (x - torch.mm(avg_emb,self.W_mean))/torch.mm(std_emb,self.W_std)
To summarize, I am projecting the LSTM output to have the same dimension as the input of my other network i.e. 840. This also corresponds to the size of the W_mean and W_std weight matrices, i.e. (840x840)
Thanks,
Gautam