[TTS] Fixing RADTTS training - removing view buffer and fixing accuracy issue

nemo/collections/tts/modules/common.py

@@ -123,15 +123,15 @@ def lstm_tensor(self, context: Tensor, lens: Tensor, enforce_sorted: bool = Fals
         seq = nn.utils.rnn.pack_padded_sequence(
             context, lens.long().cpu(), batch_first=True, enforce_sorted=enforce_sorted
         )
-        if not torch.jit.is_scripting():
-            self.bilstm.flatten_parameters()
+        # if not torch.jit.is_scripting():
+        #    self.bilstm.flatten_parameters()
         ret, _ = self.bilstm(seq)
         return nn.utils.rnn.pad_packed_sequence(ret, batch_first=True)
 
     @torch.jit.export
     def lstm_sequence(self, seq: PackedSequence) -> Tuple[Tensor, Tensor]:
-        if not torch.jit.is_scripting():
-            self.bilstm.flatten_parameters()
+        # if not torch.jit.is_scripting():
+        #    self.bilstm.flatten_parameters()
         ret, _ = self.bilstm(seq)
         return nn.utils.rnn.pad_packed_sequence(ret, batch_first=True)
 
@@ -164,10 +164,10 @@ def __init__(
     ):
         super(ConvLSTMLinear, self).__init__(n_channels, int(n_channels // 2), 1)
         self.out_dim = out_dim
+        self.convolutions = nn.ModuleList()
 
         if n_layers > 0:
             self.dropout = nn.Dropout(p=p_dropout)
-            self.convolutions = nn.ModuleList()
 
         for i in range(n_layers):
             conv_layer = ConvNorm(
@@ -219,7 +219,6 @@ def conv_to_padded_tensor(self, context: Tensor, lens: Tensor) -> Tensor:
         ret = torch.nn.utils.rnn.pad_sequence(context_embedded, batch_first=True)
         return ret
 
-    @torch.jit.export
     def masked_conv_to_sequence(self, context: Tensor, lens: Tensor, enforce_sorted: bool = False) -> PackedSequence:
         mask = get_mask_from_lengths_and_val(lens, context)
         mask = mask.unsqueeze(1)
@@ -232,24 +231,36 @@ def masked_conv_to_sequence(self, context: Tensor, lens: Tensor, enforce_sorted:
         )
         return seq
 
-    def forward(self, context: Tensor, lens: Optional[Tensor] = None) -> Tensor:
-        if lens is None:
-            for conv in self.convolutions:
-                context = self.dropout(F.relu(conv(context)))
-            context = context.transpose(1, 2)
-            context, _ = self.bilstm(context)
+    def forward(self, context: Tensor, in_lens: Optional[Tensor]) -> Tensor:
+        if in_lens is None:
+            lens = context.new_ones([context.shape[0]], dtype=torch.int) * context.shape[2]
         else:
-            # borisf : does not match ADLR (values, lengths)
-            # seq = self.masked_conv_to_sequence(context, lens, enforce_sorted=False)
-            # borisf : does match ADLR
-            seq = self.conv_to_sequence(context, lens, enforce_sorted=False)
-            context, _ = self.lstm_sequence(seq)
+            lens = in_lens
+        # borisf : does not match ADLR (values, lengths)
+        # seq = self.masked_conv_to_sequence(context, lens, enforce_sorted=False)
+        # borisf : does match ADLR
+        seq = self.conv_to_sequence(context, lens, enforce_sorted=False)
+        context, _ = self.lstm_sequence(seq)
 
         if self.dense is not None:
             context = self.dense(context).permute(0, 2, 1)
 
         return context
 
+    def script(self):
+        traced = nn.ModuleList()
+        for conv in self.convolutions:
+            if hasattr(conv, 'conv'):
+                w = conv.conv.weight
+            else:
+                w = conv.weight
+            s = w.shape
+            rand_in = torch.ones(1, s[1], s[2], dtype=w.dtype, device=w.device)
+            # , torch.ones(1, 1, s[2], dtype=w.dtype, device=w.device))
+            traced.append(torch.jit.trace_module(conv, {"forward": rand_in}))
+        self.convolutions = traced
+        return torch.jit.script(self)
+
 
 def getRadTTSEncoder(
     encoder_n_convolutions=3,

nemo/collections/tts/modules/radtts.py

@@ -784,11 +784,11 @@ def _prepare_for_export(self, **kwargs):
         PartialConv1d.forward = PartialConv1d.forward_no_cache
         self.remove_norms()
         super()._prepare_for_export(**kwargs)
-        self.encoder = torch.jit.script(self.encoder)
-        self.v_pred_module.feat_pred_fn = torch.jit.script(self.v_pred_module.feat_pred_fn)
-        self.f0_pred_module.feat_pred_fn = torch.jit.script(self.f0_pred_module.feat_pred_fn)
-        self.energy_pred_module.feat_pred_fn = torch.jit.script(self.energy_pred_module.feat_pred_fn)
-        self.dur_pred_layer.feat_pred_fn = torch.jit.script(self.dur_pred_layer.feat_pred_fn)
+        self.encoder = self.encoder.script()
+        self.v_pred_module.feat_pred_fn = self.v_pred_module.feat_pred_fn.script()
+        self.f0_pred_module.feat_pred_fn = self.f0_pred_module.feat_pred_fn.script()
+        self.energy_pred_module.feat_pred_fn = self.energy_pred_module.feat_pred_fn.script()
+        self.dur_pred_layer.feat_pred_fn = self.dur_pred_layer.feat_pred_fn.script()
 
         if self.use_context_lstm:
             self.context_lstm = torch.jit.script(self.context_lstm)

nemo/collections/tts/modules/submodules.py

@@ -30,12 +30,11 @@ def __init__(self, *args, **kwargs):
         super(PartialConv1d, self).__init__(*args, **kwargs)
         weight_maskUpdater = torch.ones(1, 1, self.kernel_size[0])
         self.register_buffer("weight_maskUpdater", weight_maskUpdater, persistent=False)
-        slide_winsize = torch.tensor(self.weight_maskUpdater.shape[1] * self.weight_maskUpdater.shape[2])
+        slide_winsize = torch.tensor(
+            self.weight_maskUpdater.shape[1] * self.weight_maskUpdater.shape[2], requires_grad=False
+        )
         self.register_buffer("slide_winsize", slide_winsize, persistent=False)
 
-        if self.bias is not None:
-            bias_view = self.bias.view(1, self.out_channels, 1)
-            self.register_buffer('bias_view', bias_view, persistent=False)
         # caching part
         self.last_size = (-1, -1, -1)
 
@@ -51,6 +50,8 @@ def calculate_mask(self, input: torch.Tensor, mask_in: Optional[torch.Tensor]):
                 mask = torch.ones(1, 1, input.shape[2], dtype=input.dtype, device=input.device)
             else:
                 mask = mask_in
+                input = torch.mul(input, mask)
+
             update_mask = F.conv1d(
                 mask,
                 self.weight_maskUpdater,
@@ -60,19 +61,22 @@ def calculate_mask(self, input: torch.Tensor, mask_in: Optional[torch.Tensor]):
                 dilation=self.dilation,
                 groups=1,
             )
-            # for mixed precision training, change 1e-8 to 1e-6
-            mask_ratio = self.slide_winsize / (update_mask + 1e-6)
+
+            update_mask_filled = torch.masked_fill(update_mask, update_mask == 0, self.slide_winsize)
+            mask_ratio = self.slide_winsize / (update_mask_filled)
+
             update_mask = torch.clamp(update_mask, 0, 1)
-            mask_ratio = torch.mul(mask_ratio.to(update_mask), update_mask)
-            return torch.mul(input, mask), mask_ratio, update_mask
+            mask_ratio = torch.mul(mask_ratio, update_mask)
+            return input, mask_ratio, update_mask
 
     def forward_aux(self, input: torch.Tensor, mask_ratio: torch.Tensor, update_mask: torch.Tensor) -> torch.Tensor:
         assert len(input.shape) == 3
 
         raw_out = self._conv_forward(input, self.weight, self.bias)
 
         if self.bias is not None:
-            output = torch.mul(raw_out - self.bias_view, mask_ratio) + self.bias_view
+            bias_view = self.bias.view(1, self.out_channels, 1)
+            output = torch.mul(raw_out - bias_view, mask_ratio) + bias_view
             output = torch.mul(output, update_mask)
         else:
             output = torch.mul(raw_out, mask_ratio)