File size: 1,302 Bytes
ad16788 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import torch
from espnet2.enh.encoder.abs_encoder import AbsEncoder
class ConvEncoder(AbsEncoder):
"""Convolutional encoder for speech enhancement and separation """
def __init__(
self,
channel: int,
kernel_size: int,
stride: int,
):
super().__init__()
self.conv1d = torch.nn.Conv1d(
1, channel, kernel_size=kernel_size, stride=stride, bias=False
)
self.stride = stride
self.kernel_size = kernel_size
self._output_dim = channel
@property
def output_dim(self) -> int:
return self._output_dim
def forward(self, input: torch.Tensor, ilens: torch.Tensor):
"""Forward.
Args:
input (torch.Tensor): mixed speech [Batch, sample]
ilens (torch.Tensor): input lengths [Batch]
Returns:
feature (torch.Tensor): mixed feature after encoder [Batch, flens, channel]
"""
assert input.dim() == 2, "Currently only support single channle input"
input = torch.unsqueeze(input, 1)
feature = self.conv1d(input)
feature = torch.nn.functional.relu(feature)
feature = feature.transpose(1, 2)
flens = (ilens - self.kernel_size) // self.stride + 1
return feature, flens
|