[wenet] support word embedding for encoder

wenet/transformer/encoder.py

@@ -31,6 +31,7 @@
 from wenet.transformer.subsampling import Conv2dSubsampling4
 from wenet.transformer.subsampling import Conv2dSubsampling6
 from wenet.transformer.subsampling import Conv2dSubsampling8
+from wenet.transformer.subsampling import EmbedinigNoSubsampling
 from wenet.transformer.subsampling import LinearNoSubsampling
 from wenet.utils.common import get_activation
 from wenet.utils.mask import make_pad_mask
@@ -104,6 +105,8 @@ def __init__(
             subsampling_class = Conv2dSubsampling6
         elif input_layer == "conv2d8":
             subsampling_class = Conv2dSubsampling8
+        elif input_layer == "embed":
+            subsampling_class = EmbedinigNoSubsampling
         else:
             raise ValueError("unknown input_layer: " + input_layer)
 

wenet/transformer/subsampling.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Modified from ESPnet(https://github.com/espnet/espnet)
-
-
 """Subsampling layer definition."""
 
 from typing import Tuple, Union
@@ -22,6 +20,7 @@
 
 
 class BaseSubsampling(torch.nn.Module):
+
     def __init__(self):
         super().__init__()
         self.right_context = 0
@@ -32,6 +31,40 @@ def position_encoding(self, offset: Union[int, torch.Tensor],
         return self.pos_enc.position_encoding(offset, size)
 
 
+class EmbedinigNoSubsampling(BaseSubsampling):
+    """Embedding input without subsampling
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        super().__init__()
+        self.embed = torch.nn.Embedding(idim, odim)
+        self.pos_enc = pos_enc_class
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Input x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: linear input tensor (#batch, time', odim),
+                where time' = time .
+            torch.Tensor: linear input mask (#batch, 1, time'),
+                where time' = time .
+
+        """
+        x = self.embed(x)
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask
+
+
 class LinearNoSubsampling(BaseSubsampling):
     """Linear transform the input without subsampling
 
@@ -41,6 +74,7 @@ class LinearNoSubsampling(BaseSubsampling):
         dropout_rate (float): Dropout rate.
 
     """
+
     def __init__(self, idim: int, odim: int, dropout_rate: float,
                  pos_enc_class: torch.nn.Module):
         """Construct an linear object."""
@@ -55,10 +89,10 @@ def __init__(self, idim: int, odim: int, dropout_rate: float,
         self.subsampling_rate = 1
 
     def forward(
-            self,
-            x: torch.Tensor,
-            x_mask: torch.Tensor,
-            offset: Union[int, torch.Tensor] = 0
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Input x.
 
@@ -87,6 +121,7 @@ class Conv2dSubsampling4(BaseSubsampling):
         dropout_rate (float): Dropout rate.
 
     """
+
     def __init__(self, idim: int, odim: int, dropout_rate: float,
                  pos_enc_class: torch.nn.Module):
         """Construct an Conv2dSubsampling4 object."""
@@ -107,10 +142,10 @@ def __init__(self, idim: int, odim: int, dropout_rate: float,
         self.right_context = 6
 
     def forward(
-            self,
-            x: torch.Tensor,
-            x_mask: torch.Tensor,
-            offset: Union[int, torch.Tensor] = 0
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Subsample x.
 
@@ -142,6 +177,7 @@ class Conv2dSubsampling6(BaseSubsampling):
         dropout_rate (float): Dropout rate.
         pos_enc (torch.nn.Module): Custom position encoding layer.
     """
+
     def __init__(self, idim: int, odim: int, dropout_rate: float,
                  pos_enc_class: torch.nn.Module):
         """Construct an Conv2dSubsampling6 object."""
@@ -160,10 +196,10 @@ def __init__(self, idim: int, odim: int, dropout_rate: float,
         self.right_context = 10
 
     def forward(
-            self,
-            x: torch.Tensor,
-            x_mask: torch.Tensor,
-            offset: Union[int, torch.Tensor] = 0
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Subsample x.
         Args:
@@ -194,6 +230,7 @@ class Conv2dSubsampling8(BaseSubsampling):
         dropout_rate (float): Dropout rate.
 
     """
+
     def __init__(self, idim: int, odim: int, dropout_rate: float,
                  pos_enc_class: torch.nn.Module):
         """Construct an Conv2dSubsampling8 object."""
@@ -214,10 +251,10 @@ def __init__(self, idim: int, odim: int, dropout_rate: float,
         self.right_context = 14
 
     def forward(
-            self,
-            x: torch.Tensor,
-            x_mask: torch.Tensor,
-            offset: Union[int, torch.Tensor] = 0
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Subsample x.