CausalConv1d module with customized initialization.
| 10 | |
| 11 | |
| 12 | class CausalConv1d(torch.nn.Module): |
| 13 | """CausalConv1d module with customized initialization.""" |
| 14 | |
| 15 | def __init__(self, in_channels, out_channels, kernel_size, |
| 16 | dilation=1, bias=True, pad="ConstantPad1d", pad_params={"value": 0.0}): |
| 17 | """Initialize CausalConv1d module.""" |
| 18 | super(CausalConv1d, self).__init__() |
| 19 | self.pad = getattr(torch.nn, pad)((kernel_size - 1) * dilation, **pad_params) |
| 20 | self.conv = torch.nn.Conv1d(in_channels, out_channels, kernel_size, |
| 21 | dilation=dilation, bias=bias) |
| 22 | |
| 23 | def forward(self, x): |
| 24 | """Calculate forward propagation. |
| 25 | |
| 26 | Args: |
| 27 | x (Tensor): Input tensor (B, in_channels, T). |
| 28 | |
| 29 | Returns: |
| 30 | Tensor: Output tensor (B, out_channels, T). |
| 31 | |
| 32 | """ |
| 33 | return self.conv(self.pad(x))[:, :, :x.size(2)] |
| 34 | |
| 35 | |
| 36 | class CausalConvTranspose1d(torch.nn.Module): |