WeNet和ESPnet中下采样模块(Conv2dSubsampling)

关于WeNet和ESPnet两个工具下采样模块都是相同的操作,

首先将输入序列扩充一个维度(因为要使用二维卷积),

然后通过两个二维卷积,其中第一个卷积的输入通道为“1”,输出通道为odim(ESPnet中默认为256,WeNet默认为512),卷积核大小为3x3。

第二个卷积输入通道是odim,输出通道也是odim,卷积核大小为3x3。

通过两个卷积后,输入特征的时间维度T和特征维度dim分别被压缩一半,然后把时间维度转换到维度1位置,再将最后两个维度融合到一起,并且通过线性变换恢复到输入维度上。

代码如下(ESPnert):

class Conv2dSubsampling(torch.nn.Module):
    """Convolutional 2D subsampling (to 1/4 length).

    Args:
        idim (int): Input dimension.
        odim (int): Output dimension.
        dropout_rate (float): Dropout rate.
        pos_enc (torch.nn.Module): Custom position encoding layer.

    """

    def __init__(self, idim, odim, dropout_rate, pos_enc=None):
        """Construct an Conv2dSubsampling object."""
        super(Conv2dSubsampling, self).__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1, odim, 3, 2),
            torch.nn.ReLU(),
            torch.nn.Conv2d(odim, odim, 3, 2),
            torch.nn.ReLU(),
        )
        self.out = torch.nn.Sequential(
            torch.nn.Linear(odim * (((idim - 1) // 2 - 1) // 2), odim), #((idim - 1) // 2 - 1) // 2=19
            pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate),
        )

    def forward(self, x, x_mask):
        """Subsample x.

        Args:
            x (torch.Tensor): Input tensor (#batch, time, idim).
            x_mask (torch.Tensor): Input mask (#batch, 1, time).

        Returns:
            torch.Tensor: Subsampled tensor (#batch, time', odim),
                where time' = time // 4.
            torch.Tensor: Subsampled mask (#batch, 1, time'),
                where time' = time // 4.

        """
        #x.size=[98,528,80]  80是通过前端后的特征维度,每个批量进来都是80维
        x = x.unsqueeze(1)  # (b, c, t, f) [98,1,528, 80]
        x = self.conv(x)  #[98, 256, 131, 19]   #256是输出维度,131是压缩4倍后句子的长度,19也是每个批量都是保持一样的
        b, c, t, f = x.size()
        x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) #[98, 131, 256]
        if x_mask is None:
            return x, None
        return x, x_mask[:, :, :-2:2][:, :, :-2:2]

    def __getitem__(self, key):
        """Get item.

        When reset_parameters() is called, if use_scaled_pos_enc is used,
            return the positioning encoding.

        """
        if key != -1:
            raise NotImplementedError("Support only `-1` (for `reset_parameters`).")
        return self.out[key]

 

posted @ 2022-11-24 22:07  Uriel-w  阅读(305)  评论(0编辑  收藏  举报