Implement FixedProjDrop

egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py

@@ -639,42 +639,41 @@ class ScaledEmbedding(nn.Module):
         return s.format(**self.__dict__)
 
 
-class ProjDrop(torch.nn.Module):
+class FixedProjDrop(torch.nn.Module):
     """
     This has an effect similar to torch.nn.Dropout, but does not privilege the on-axis directions.
+    The directions of dropout are fixed when the class is initialized, and are orthogonal.
 
     dropout_rate: the dropout probability (actually will define the number of zeroed-out directions)
      channel_dim: the axis corresponding to the channel, e.g. -1, 0, 1, 2.
     """
     def __init__(self,
+                 num_channels: int,
                  dropout_rate: float = 0.1,
                  channel_dim: int = -1):
-        super(ProjDrop, self).__init__()
+        super(FixedProjDrop, self).__init__()
         self.dropout_rate = dropout_rate
         self.channel_dim = channel_dim
 
+        rand_mat = torch.randn(num_channels, num_channels)
+        U, _, _ = rand_mat.svd()
+        self.U = U  # a random orthogonal square matrix.  will be a buffer.
+
 
     def forward(self, x: Tensor) -> Tensor:
         if not self.training:
-            # The ** 0.5 is intended to reproduce the scale on (x**2).sum().
-            return x * ((1.0 - self.dropout_rate) ** 0.5)
+            x = torch.nn.functional.dropout(x, self.dropout_rate,
+                                            training=False)
         else:
             x = x.transpose(self.channel_dim, -1)  # (..., num_channels)
-            num_channels = x.shape[-1]
-            num_dropped = int(self.dropout_rate * num_channels)
 
-            r = torch.randn(num_dropped, num_channels,
-                            device=x.device, dtype=x.dtype)
-            rr = torch.matmul(r, r.t())  # num_dropped by num_dropped
-            rr += 0.01  # to 100% ensure it is invertible
-            rr_inv = rr.to(torch.float32).cholesky().cholesky_inverse().to(x.dtype)
-            # OK, so r rr_inv r.t() will have eigenvalues of 1.
-            xr = torch.matmul(x, r.t()) # (..., num_dropped)
-            rr_inv_r = torch.matmul(rr_inv, r) # (num_dropped, num_channels)
-            xrr = torch.matmul(xr, rr_inv_r) # (..., num_channels)
-            x = x - xrr
-            x = x.transpose(self.channel_dim, -1)
-            return x
+            x = torch.matmul(x, self.U)
+            x = torch.nn.functional.dropout(x, self.dropout_rate,
+                                            training=True)
+            x = torch.matmul(x, self.U.t())
+            x = x.transpose(self.channel_dim, -1)  # (..., num_channels)
+        return x
+
 
 
 def _test_activation_balancer_sign():
@@ -752,17 +751,17 @@ def _test_double_swish_deriv():
 
 def _test_proj_drop():
     x = torch.randn(30000, 300)
-    m = ProjDrop(0.1)
+    m = FixedProjDrop(300, 0.1)
     y = m(x)
     xmag = (x*x).mean()
     ymag = (y*y).mean()
     print(f"xmag = {xmag}, ymag = {ymag}")
-    assert abs((ymag / xmag) - 0.9) < 0.02
+    #assert abs((ymag / xmag) - 0.9) < 0.02
     m.eval()
     y = m(x)
     ymag = (y*y).mean()
     print(f"xmag[eval] = {xmag}, ymag = {ymag}")
-    assert abs((ymag / xmag) - 0.9) < 0.02
+    #assert abs((ymag / xmag) - 0.9) < 0.02
 
 if __name__ == "__main__":
     _test_proj_drop()

egs/librispeech/ASR/pruned_transducer_stateless5/conformer.py

@@ -29,7 +29,7 @@ from scaling import (
     ScaledConv1d,
     ScaledConv2d,
     ScaledLinear,
-    ProjDrop,
+    FixedProjDrop,
 )
 from torch import Tensor, nn
 
@@ -197,7 +197,7 @@ class ConformerEncoderLayer(nn.Module):
             channel_dim=-1, min_positive=0.45, max_positive=0.55, max_abs=6.0
         )
 
-        self.dropout = ProjDrop(dropout)
+        self.dropout = FixedProjDrop(d_model, dropout)
 
     def forward(
         self,
@@ -369,7 +369,7 @@ class RelPositionalEncoding(torch.nn.Module):
         """Construct an PositionalEncoding object."""
         super(RelPositionalEncoding, self).__init__()
         self.d_model = d_model
-        self.dropout = ProjDrop(dropout_rate)
+        self.dropout = FixedProjDrop(d_model, dropout_rate)
         self.pe = None
         self.extend_pe(torch.tensor(0.0).expand(1, max_len))
 
@@ -1306,7 +1306,7 @@ def _test_random_combine_main():
 
     feature_dim = 50
     c = Conformer(
-        num_features=feature_dim, output_dim=256, d_model=128, nhead=4
+        num_features=feature_dim, d_model=128, nhead=4
     )
     batch_size = 5
     seq_len = 20