I have implemented a custom rnn layer, it is a little different with standard rnn layer. I found the official LSTM/GRU’s implementation doesn’t use for loop, my question is the RNN’s backpropagation is BPTT, if I use a for loop in model, could it get the right BPTT grad?
class MyRNN(nn.Module):
def __ini__(self):
super(MyRNN, self).__init__()
self.hidden_dim = 8
self.embedding_dim = 8
self.embedding = nn.Embedding(n_words, self.embedding_dim)
self.h = torch.randn(batch_size, self.hidden_dim)
self.linear_x = nn.Linear(self.embedding_dim, self.hidden_dim)
self.linear_h = nn.Linear(self.hidden_dim, self.hidden_dim)
def forward(self, x):
x = self.embedding(x)
h = self.h
for i in range(x.shape[1]): # steps
linear_x = self.linear_x(x[:, i, :])
h = self.linear_h(h)
h = F.tanh( linearx + h)
return h
