| 806 | |
| 807 | |
| 808 | class CausalConv3d(Module): |
| 809 | @beartype |
| 810 | def __init__( |
| 811 | self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], pad_mode="constant", **kwargs |
| 812 | ): |
| 813 | super().__init__() |
| 814 | kernel_size = cast_tuple(kernel_size, 3) |
| 815 | |
| 816 | time_kernel_size, height_kernel_size, width_kernel_size = kernel_size |
| 817 | |
| 818 | assert is_odd(height_kernel_size) and is_odd(width_kernel_size) |
| 819 | |
| 820 | dilation = kwargs.pop("dilation", 1) |
| 821 | stride = kwargs.pop("stride", 1) |
| 822 | |
| 823 | self.pad_mode = pad_mode |
| 824 | time_pad = dilation * (time_kernel_size - 1) + (1 - stride) |
| 825 | height_pad = height_kernel_size // 2 |
| 826 | width_pad = width_kernel_size // 2 |
| 827 | |
| 828 | self.time_pad = time_pad |
| 829 | self.time_causal_padding = (width_pad, width_pad, height_pad, height_pad, time_pad, 0) |
| 830 | |
| 831 | stride = (stride, 1, 1) |
| 832 | dilation = (dilation, 1, 1) |
| 833 | self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) |
| 834 | |
| 835 | def forward(self, x): |
| 836 | pad_mode = self.pad_mode if self.time_pad < x.shape[2] else "constant" |
| 837 | |
| 838 | x = F.pad(x, self.time_causal_padding, mode=pad_mode) |
| 839 | return self.conv(x) |
| 840 | |
| 841 | |
| 842 | @beartype |
no outgoing calls
no test coverage detected