Hi everyone,
I need to implement the Gated Recurrent Unit by myself because I will need to modify some terms in its equations. But, firstly I would like to implement the original GRU to be sure that everything works correctly. This is my implementation, I did a simple test to compare its results with the ones of the Pytorch implementation, but they are different.
Where am I going wrong?
import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
import math
class GRU(nn.Module):
def __init__(self, input_size, hidden_size):
super(GRU, self).__init__()
self.Wx = Parameter(torch.randn(input_size, hidden_size * 3)).div(math.sqrt(input_size))
self.bx = Parameter(torch.zeros(hidden_size * 3))
self.Wh = Parameter(torch.randn(hidden_size, hidden_size * 3)).div(math.sqrt(hidden_size))
self.bh = Parameter(torch.zeros(hidden_size * 3))
def step(self, x, prev_h):
batch_size, hidden_size = prev_h.shape
ax = x.mm(self.Wx) + self.bx
ax = ax.reshape(batch_size, 3, hidden_size)
ah = prev_h.mm(self.Wh) + self.bh
ah = ah.reshape(batch_size, 3, hidden_size)
r = torch.sigmoid(ax[:, 0, :] + ah[:, 0, :])
z = torch.sigmoid(ax[:, 1, :] + ah[:, 1, :])
n = torch.tanh(ax[:, 2, :] + r * ah[:, 2, :])
next_h = (1 - z) * n + z * prev_h
return next_h
def forward(self, x, h0):
batch_size, steps, num_features = x.shape
hidden_size = h0.shape[1]
hts = torch.zeros(batch_size, steps, hidden_size).to(dtype=x.dtype)
h_t = h0
for step_t in range(steps):
x_t = x[:, step_t]
h_t = self.step(x_t, h_t)
hts[:, step_t] = h_t
return hts
Thank you 