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]