add l2 norm for embeddings, +minor modification to make it more consistent with tf version

Author: pmixer

main.py

@@ -23,7 +23,7 @@ def str2bool(s):
 parser.add_argument('--num_blocks', default=2, type=int)
 parser.add_argument('--num_epochs', default=201, type=int)
 parser.add_argument('--num_heads', default=1, type=int)
-parser.add_argument('--dropout_rate', default=0.5, type=float)
+parser.add_argument('--dropout_rate', default=0.2, type=float)
 parser.add_argument('--l2_emb', default=0.0, type=float)
 parser.add_argument('--device', default='cpu', type=str)
 parser.add_argument('--inference_only', default=False, type=str2bool)
@@ -91,9 +91,14 @@ def str2bool(s):
         indices = np.where(pos != 0)
         loss = bce_criterion(pos_logits[indices], pos_labels[indices])
         loss += bce_criterion(neg_logits[indices], neg_labels[indices])
+        for param in model.item_emb.parameters(): loss += args.l2_emb * torch.norm(param)
+        for param in model.abs_pos_K_emb.parameters(): loss += args.l2_emb * torch.norm(param)
+        for param in model.abs_pos_V_emb.parameters(): loss += args.l2_emb * torch.norm(param)
+        for param in model.time_matrix_K_emb.parameters(): loss += args.l2_emb * torch.norm(param)
+        for param in model.time_matrix_V_emb.parameters(): loss += args.l2_emb * torch.norm(param)
         loss.backward()
         adam_optimizer.step()
-        # print("loss in epoch {} iteration {}: {}".format(epoch, step, loss.item())) # expected 0.4~0.6 after init few epochs
+        print("loss in epoch {} iteration {}: {}".format(epoch, step, loss.item())) # expected 0.4~0.6 after init few epochs
 
     if epoch % 20 == 0:
         model.eval()

model.py

@@ -64,19 +64,20 @@ def forward(self, queries, keys, time_mask, attn_mask, time_matrix_K, time_matri
 
         time_mask = time_mask.unsqueeze(-1).expand(attn_weights.shape[0], -1, attn_weights.shape[-1])
         attn_mask = attn_mask.unsqueeze(0).expand(attn_weights.shape[0], -1, -1)
-        paddings = torch.ones(attn_weights.shape) *  -1e23 # float('-inf')
+        paddings = torch.ones(attn_weights.shape) *  (-2**32+1) # -1e23 # float('-inf')
         paddings = paddings.to(self.dev)
         attn_weights = torch.where(time_mask, paddings, attn_weights) # True:pick padding
         attn_weights = torch.where(attn_mask, paddings, attn_weights) # enforcing causality
 
         attn_weights = self.softmax(attn_weights) # code as below invalids pytorch backward rules
+        # attn_weights = torch.where(time_mask, paddings, attn_weights) # weird query mask in tf impl
         # https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/4
         # attn_weights[attn_weights != attn_weights] = 0 # rm nan for -inf into softmax case
         attn_weights = self.dropout(attn_weights)
 
         outputs = attn_weights.matmul(V_)
         outputs += attn_weights.matmul(abs_pos_V_)
-        outputs += attn_weights.unsqueeze(-2).matmul(time_matrix_V_).reshape(outputs.shape)
+        outputs += attn_weights.unsqueeze(2).matmul(time_matrix_V_).reshape(outputs.shape).squeeze(2)
 
         # (num_head * N, T, C / num_head) -> (N, T, C)
         outputs = torch.cat(torch.split(outputs, Q.shape[0], dim=0), dim=2) # div batch_size