update for training

nobody132 · Jun 20, 2019 · b5209eb · b5209eb
1 parent 82371b5
commit b5209eb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/models/conv.py b/models/conv.py
@@ -10,7 +10,7 @@ def __init__(self, conv, p):
         super().__init__()
         self.conv = conv
         nn.init.kaiming_normal_(self.conv.weight)
-        weight_norm(self.conv)
+        self.conv = weight_norm(self.conv)
         self.act = nn.GLU(1)
         self.dropout = nn.Dropout(p, inplace=True)
 
@@ -50,7 +50,9 @@ def forward(self, x, lens):  # -> B * V * T
         x = self.cnn(x)
         for module in self.modules():
             if type(module) == nn.modules.Conv1d:
-                lens = (lens - module.kernel_size[0]) / module.stride[0] + 1
+                lens = (
+                    lens - module.kernel_size[0] + 2 * module.padding[0]
+                ) // module.stride[0] + 1
         return x, lens
 
     def predict(self, path):

diff --git a/train.py b/train.py
@@ -75,7 +75,7 @@ def train(
         writer.add_scalar("loss/epoch", epoch_loss, epoch)
         writer.add_scalar("cer/epoch", cer, epoch)
         print("Epoch {}: Loss= {}, CER = {}".format(epoch, epoch_loss, cer))
-        torch.save(model, "models/v5/casr_v5_1_epoch_{}.pth".format(epoch))
+        torch.save(model, "pretrained/model_{}.pth".format(epoch))
 
 
 def get_lr(optimizer):