| 225 | |
| 226 | class CausalConv(nn.Module): |
| 227 | def __init__(self, |
| 228 | chan_in, |
| 229 | chan_out, |
| 230 | kernel_size, |
| 231 | **kwargs |
| 232 | ): |
| 233 | super().__init__() |
| 234 | |
| 235 | if isinstance(kernel_size, int): |
| 236 | kernel_size = kernel_size if isinstance(kernel_size, tuple) else ((kernel_size,) * 3) |
| 237 | time_kernel_size, height_kernel_size, width_kernel_size = kernel_size |
| 238 | |
| 239 | self.dilation = kwargs.pop('dilation', 1) |
| 240 | self.stride = kwargs.pop('stride', 1) |
| 241 | if isinstance(self.stride, int): |
| 242 | self.stride = (self.stride, 1, 1) |
| 243 | time_pad = self.dilation * (time_kernel_size - 1) + max((1 - self.stride[0]), 0) |
| 244 | height_pad = height_kernel_size // 2 |
| 245 | width_pad = width_kernel_size // 2 |
| 246 | self.time_causal_padding = (width_pad, width_pad, height_pad, height_pad, time_pad, 0) |
| 247 | self.time_uncausal_padding = (width_pad, width_pad, height_pad, height_pad, 0, 0) |
| 248 | |
| 249 | self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs) |
| 250 | self.is_first_run = True |
| 251 | |
| 252 | def forward(self, x, is_init=True, residual=None): |
| 253 | x = nn.functional.pad(x, |