reinsert norm to spell, changed loss as mean

Author: jinserk

asr/models/las/network.py

@@ -200,7 +200,7 @@ def __init__(self, listen_vec_size, label_vec_size, max_seq_lens=256, sos=None,
         self.rnn_num_layers = rnn_num_layers
         self.rnns = rnn_type(input_size=(Hy + Hc), hidden_size=Hs, num_layers=rnn_num_layers,
                              bias=True, bidirectional=False, batch_first=True)
-        #self.norm = nn.LayerNorm(Hs, elementwise_affine=False)
+        self.norm = nn.LayerNorm(Hs, elementwise_affine=False)
 
         self.attention = Attention(state_vec_size=Hs, listen_vec_size=Hc,
                                    apply_proj=apply_attend_proj, proj_hidden_size=proj_hidden_size,
@@ -242,7 +242,7 @@ def forward(self, h, x_seq_lens, y=None, y_seq_lens=None):
 
         for t in range(self.max_seq_lens):
             s, hidden = self.rnns(x, hidden)
-            #s = self.norm(s)
+            s = self.norm(s)
             c, a = self.attention(s, h, in_mask)
             y_hat = self.chardist(torch.cat([s, c], dim=-1))
             y_hat = self.softmax(y_hat)
@@ -407,7 +407,7 @@ def _eval_forward(self, x, x_seq_lens):
         h = self.listen(x, x_seq_lens)
         # spell
         y_hats, y_hats_seq_lens, _ = self.spell(h, x_seq_lens)
-        y_hats_seq_lens[y_hats.seq_lens.ne(self.spell.max_seq_lens)].sub_(self.spell.num_eos)
+        y_hats_seq_lens[y_hats_seq_lens.ne(self.spell.max_seq_lens)].sub_(self.spell.num_eos)
 
         # return with seq lens without sos and eos
         y_hats = self.log(y_hats[:, :, :-2])

asr/models/las/train.py

@@ -27,7 +27,7 @@ class LASTrainer(NonSplitTrainer):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self.loss = nn.NLLLoss(reduction='none', ignore_index=self.model.blk)
+        self.loss = nn.NLLLoss(reduction='mean', ignore_index=self.model.blk)
 
         self.tfr_scheduler = TFRScheduler(self.model, ranges=(0.9, 0.0), warm_up=0, epochs=9, restart=True)
         #self.tfr_scheduler.step(9)
@@ -68,7 +68,8 @@ def unit_train(self, data):
                 ys_hat = ys_hat.float()
             if self.use_cuda:
                 ys_lens = ys_lens.cuda()
-            loss = self.loss(ys_hat.transpose(1, 2), ys.long()).sum(dim=-1).div(ys_lens.float()).mean()
+            loss = self.loss(ys_hat.transpose(1, 2), ys.long())
+            #loss = self.loss(ys_hat.transpose(1, 2), ys.long()).sum(dim=-1).div(ys_lens.float()).mean()
             #if ys_hat_lens is None:
             #    logger.debug("the batch includes a data with label_lens > max_seq_lens: ignore the entire batch")
             #    loss.mul_(0)