From f4941c525f9eff0ed5c56e31e3b8f14564211f4d Mon Sep 17 00:00:00 2001 From: Kevin Date: Mon, 6 Aug 2018 06:25:12 -0700 Subject: [PATCH] shared weights fix --- models/nalu.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/models/nalu.py b/models/nalu.py index 65015e4..2b12d06 100644 --- a/models/nalu.py +++ b/models/nalu.py @@ -25,19 +25,17 @@ class NeuralArithmeticLogicUnitCell(nn.Module): self.eps = 1e-10 self.G = Parameter(torch.Tensor(out_dim, in_dim)) - self.W = Parameter(torch.Tensor(out_dim, in_dim)) - self.register_parameter('bias', None) self.nac = NeuralAccumulatorCell(in_dim, out_dim) + self.register_parameter('bias', None) init.kaiming_uniform_(self.G, a=math.sqrt(5)) - init.kaiming_uniform_(self.W, a=math.sqrt(5)) def forward(self, input): a = self.nac(input) g = F.sigmoid(F.linear(input, self.G, self.bias)) add_sub = g * a log_input = torch.log(torch.abs(input) + self.eps) - m = torch.exp(F.linear(log_input, self.W, self.bias)) + m = torch.exp(self.nac(log_input)) mul_div = (1 - g) * m y = add_sub + mul_div return y