update

toolbox/torchaudio/models/spectrum_dfnet/modeling_spectrum_dfnet.py

@@ -802,30 +802,31 @@ class SpectrumDfNet(nn.Module):
     def forward(self,
                 spec_complex: torch.Tensor,
                 ):
-        feat_power = torch.square(torch.abs(spec_complex))
-        feat_power = feat_power.unsqueeze(1).permute(0, 1, 3, 2)
-        # feat_power shape: [batch_size, spec_bins, time_steps]
-        # feat_power shape: [batch_size, 1, spec_bins, time_steps]
-        # feat_power shape: [batch_size, 1, time_steps, spec_bins]
-        feat_power = feat_power.detach()
-
-        # spec shape: [batch_size, spec_bins, time_steps]
-        feat_spec = torch.view_as_real(spec_complex)
-        # spec shape: [batch_size, spec_bins, time_steps, 2]
-        feat_spec = feat_spec.permute(0, 3, 2, 1)
-        # feat_spec shape: [batch_size, 2, time_steps, spec_bins]
-        feat_spec = feat_spec[..., :self.df_decoder.df_bins]
-        # feat_spec shape: [batch_size, 2, time_steps, df_bins]
-        feat_spec = feat_spec.detach()
-
-        # spec shape: [batch_size, spec_bins, time_steps]
-        spec = torch.unsqueeze(spec_complex, dim=1)
-        # spec shape: [batch_size, 1, spec_bins, time_steps]
-        spec = spec.permute(0, 1, 3, 2)
-        # spec shape: [batch_size, 1, time_steps, spec_bins]
-        spec = torch.view_as_real(spec)
-        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
-        spec = spec.detach()
+        with torch.no_grad():
+            feat_power = torch.square(torch.abs(spec_complex))
+            feat_power = feat_power.unsqueeze(1).permute(0, 1, 3, 2)
+            # feat_power shape: [batch_size, spec_bins, time_steps]
+            # feat_power shape: [batch_size, 1, spec_bins, time_steps]
+            # feat_power shape: [batch_size, 1, time_steps, spec_bins]
+            feat_power = feat_power.detach()
+
+            # spec shape: [batch_size, spec_bins, time_steps]
+            feat_spec = torch.view_as_real(spec_complex)
+            # spec shape: [batch_size, spec_bins, time_steps, 2]
+            feat_spec = feat_spec.permute(0, 3, 2, 1)
+            # feat_spec shape: [batch_size, 2, time_steps, spec_bins]
+            feat_spec = feat_spec[..., :self.df_decoder.df_bins]
+            # feat_spec shape: [batch_size, 2, time_steps, df_bins]
+            feat_spec = feat_spec.detach()
+
+        # # spec shape: [batch_size, spec_bins, time_steps]
+        # spec = torch.unsqueeze(spec_complex, dim=1)
+        # # spec shape: [batch_size, 1, spec_bins, time_steps]
+        # spec = spec.permute(0, 1, 3, 2)
+        # # spec shape: [batch_size, 1, time_steps, spec_bins]
+        # spec = torch.view_as_real(spec)
+        # # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
+        # spec = spec.detach()
 
         e0, e1, e2, e3, emb, c0, lsnr, h = self.encoder.forward(feat_power, feat_spec)
 
@@ -834,24 +835,24 @@ class SpectrumDfNet(nn.Module):
         if torch.any(mask > 1) or torch.any(mask < 0):
             raise AssertionError
 
-        spec_m = self.mask.forward(spec, mask)
-
-        # lsnr shape: [batch_size, time_steps, 1]
-        lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
-        # lsnr shape: [batch_size, 1, time_steps]
-
-        df_coefs = self.df_decoder.forward(emb, c0)
-        df_coefs = self.df_out_transform(df_coefs)
-        # df_coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
-
-        spec_e = self.df_op.forward(spec.clone(), df_coefs)
-        # spec_e shape: [batch_size, 1, time_steps, spec_bins, 2]
-
-        spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]
-
-        spec_e = torch.squeeze(spec_e, dim=1)
-        spec_e = spec_e.permute(0, 2, 1, 3)
-        # spec_e shape: [batch_size, spec_bins, time_steps, 2]
+        # spec_m = self.mask.forward(spec, mask)
+        #
+        # # lsnr shape: [batch_size, time_steps, 1]
+        # lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
+        # # lsnr shape: [batch_size, 1, time_steps]
+        #
+        # df_coefs = self.df_decoder.forward(emb, c0)
+        # df_coefs = self.df_out_transform(df_coefs)
+        # # df_coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
+        #
+        # spec_e = self.df_op.forward(spec.clone(), df_coefs)
+        # # spec_e shape: [batch_size, 1, time_steps, spec_bins, 2]
+        #
+        # spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]
+        #
+        # spec_e = torch.squeeze(spec_e, dim=1)
+        # spec_e = spec_e.permute(0, 2, 1, 3)
+        # # spec_e shape: [batch_size, spec_bins, time_steps, 2]
 
         mask = torch.squeeze(mask, dim=1)
         mask = mask.permute(0, 2, 1)