def forward(self, x):
# Applying Layer/Batch Norm
if bool(self.minimalgru_use_laynorm_inp):
x = self.ln0((x))
if bool(self.minimalgru_use_batchnorm_inp):
x_bn = self.bn0(x.view(x.shape[0] * x.shape[1], x.shape[2]))
x = x_bn.view(x.shape[0], x.shape[1], x.shape[2])
for i in range(self.N_minimalgru_lay):
# Initial state and concatenation
if self.bidir:
h_init = torch.zeros(2 * x.shape[1], self.minimalgru_lay[i])
x = torch.cat([x, flip(x, 0)], 1)
else:
h_init = torch.zeros(x.shape[1], self.minimalgru_lay[i])
# Drop mask initilization (same mask for all time steps)
if self.test_flag == False:
drop_mask = torch.bernoulli(
torch.Tensor(
h_init.shape[0],
h_init.shape[1]).fill_(1 - self.minimalgru_drop[i]))
else:
drop_mask = torch.FloatTensor([1 - self.minimalgru_drop[i]])
if self.use_cuda:
h_init = h_init.cuda()
drop_mask = drop_mask.cuda()
# Feed-forward affine transformations (all steps in parallel)
wh_out = self.wh[i](x)
wz_out = self.wz[i](x)
# Apply batch norm if needed (all steos in parallel)
if self.minimalgru_use_batchnorm[i]:
wh_out_bn = self.bn_wh[i](wh_out.view(
wh_out.shape[0] * wh_out.shape[1], wh_out.shape[2]))
wh_out = wh_out_bn.view(wh_out.shape[0], wh_out.shape[1],
wh_out.shape[2])
wz_out_bn = self.bn_wz[i](wz_out.view(
wz_out.shape[0] * wz_out.shape[1], wz_out.shape[2]))
wz_out = wz_out_bn.view(wz_out.shape[0], wz_out.shape[1],
wz_out.shape[2])
# Processing time steps
hiddens = []
ht = h_init
for k in range(x.shape[0]):
# minimalgru equation
zt = torch.sigmoid(wz_out[k] + self.uz[i](ht))
at = wh_out[k] + self.uh[i](zt * ht)
hcand = self.act[i](at) * drop_mask
ht = (zt * ht + (1 - zt) * hcand)
if self.minimalgru_use_laynorm[i]:
ht = self.ln[i](ht)
hiddens.append(ht)
# Stacking hidden states
h = torch.stack(hiddens)
# Bidirectional concatenations
if self.bidir:
h_f = h[:, 0:int(x.shape[1] / 2)]
h_b = flip(h[:, int(x.shape[1] / 2):x.shape[1]].contiguous(),
0)
h = torch.cat([h_f, h_b], 2)
# Setup x for the next hidden layer
x = h
return x
def forward(self, x):
# Applying Layer/Batch Norm
if bool(self.lstm_use_laynorm_inp):
x = self.ln0((x))
if bool(self.lstm_use_batchnorm_inp):
x_bn = self.bn0(x.view(x.shape[0] * x.shape[1], x.shape[2]))
x = x_bn.view(x.shape[0], x.shape[1], x.shape[2])
for i in range(self.N_lstm_lay):
# Initial state and concatenation
if self.bidir:
h_init = torch.zeros(2 * x.shape[1], self.lstm_lay[i])
x = torch.cat([x, flip(x, 0)], 1)
else:
h_init = torch.zeros(x.shape[1], self.lstm_lay[i])
# Drop mask initilization (same mask for all time steps)
if self.test_flag == False:
drop_mask = torch.bernoulli(
torch.Tensor(h_init.shape[0],
h_init.shape[1]).fill_(1 - self.lstm_drop[i]))
else:
drop_mask = torch.FloatTensor([1 - self.lstm_drop[i]])
h_init = h_init.to(x.device)
drop_mask = drop_mask.to(x.device)
# Feed-forward affine transformations (all steps in parallel)
wfx_out = self.wfx[i](x)
wix_out = self.wix[i](x)
wox_out = self.wox[i](x)
wcx_out = self.wcx[i](x)
# Apply batch norm if needed (all steos in parallel)
if self.lstm_use_batchnorm[i]:
wfx_out_bn = self.bn_wfx[i](wfx_out.view(
wfx_out.shape[0] * wfx_out.shape[1], wfx_out.shape[2]))
wfx_out = wfx_out_bn.view(wfx_out.shape[0], wfx_out.shape[1],
wfx_out.shape[2])
wix_out_bn = self.bn_wix[i](wix_out.view(
wix_out.shape[0] * wix_out.shape[1], wix_out.shape[2]))
wix_out = wix_out_bn.view(wix_out.shape[0], wix_out.shape[1],
wix_out.shape[2])
wox_out_bn = self.bn_wox[i](wox_out.view(
wox_out.shape[0] * wox_out.shape[1], wox_out.shape[2]))
wox_out = wox_out_bn.view(wox_out.shape[0], wox_out.shape[1],
wox_out.shape[2])
wcx_out_bn = self.bn_wcx[i](wcx_out.view(
wcx_out.shape[0] * wcx_out.shape[1], wcx_out.shape[2]))
wcx_out = wcx_out_bn.view(wcx_out.shape[0], wcx_out.shape[1],
wcx_out.shape[2])
# Processing time steps
hiddens = []
ct = h_init
ht = h_init
for k in range(x.shape[0]):
# LSTM equations
ft = torch.sigmoid(wfx_out[k] + self.ufh[i](ht))
it = torch.sigmoid(wix_out[k] + self.uih[i](ht))
ot = torch.sigmoid(wox_out[k] + self.uoh[i](ht))
ct = it * self.act[i](wcx_out[k] + self.uch[i]
(ht)) * drop_mask + ft * ct
ht = ot * self.act[i](ct)
if self.lstm_use_laynorm[i]:
ht = self.ln[i](ht)
hiddens.append(ht)
# Stacking hidden states
h = torch.stack(hiddens)
# Bidirectional concatenations
if self.bidir:
h_f = h[:, 0:int(x.shape[1] / 2)]
h_b = flip(h[:, int(x.shape[1] / 2):x.shape[1]].contiguous(),
0)
h = torch.cat([h_f, h_b], 2)
# Setup x for the next hidden layer
x = h
return x