| |
|
|
| |
| |
|
|
| """Causal convolusion layer modules.""" |
|
|
|
|
| import torch |
|
|
|
|
| class CausalConv1d(torch.nn.Module): |
| """CausalConv1d module with customized initialization.""" |
|
|
| def __init__(self, in_channels, out_channels, kernel_size, |
| dilation=1, bias=True, pad="ConstantPad1d", pad_params={"value": 0.0}): |
| """Initialize CausalConv1d module.""" |
| super(CausalConv1d, self).__init__() |
| self.pad = getattr(torch.nn, pad)((kernel_size - 1) * dilation, **pad_params) |
| self.conv = torch.nn.Conv1d(in_channels, out_channels, kernel_size, |
| dilation=dilation, bias=bias) |
|
|
| def forward(self, x): |
| """Calculate forward propagation. |
| |
| Args: |
| x (Tensor): Input tensor (B, in_channels, T). |
| |
| Returns: |
| Tensor: Output tensor (B, out_channels, T). |
| |
| """ |
| return self.conv(self.pad(x))[:, :, :x.size(2)] |
|
|
|
|
| class CausalConvTranspose1d(torch.nn.Module): |
| """CausalConvTranspose1d module with customized initialization.""" |
|
|
| def __init__(self, in_channels, out_channels, kernel_size, stride, bias=True): |
| """Initialize CausalConvTranspose1d module.""" |
| super(CausalConvTranspose1d, self).__init__() |
| self.deconv = torch.nn.ConvTranspose1d( |
| in_channels, out_channels, kernel_size, stride, bias=bias) |
| self.stride = stride |
|
|
| def forward(self, x): |
| """Calculate forward propagation. |
| |
| Args: |
| x (Tensor): Input tensor (B, in_channels, T_in). |
| |
| Returns: |
| Tensor: Output tensor (B, out_channels, T_out). |
| |
| """ |
| return self.deconv(x)[:, :, :-self.stride] |
|
|
