| 167 | |
| 168 | |
| 169 | class CachedConv3d(torch.nn.Conv3d): |
| 170 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0): |
| 171 | super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding) |
| 172 | self.cached_tensor = None |
| 173 | |
| 174 | |
| 175 | def clear_cache(self): |
| 176 | self.cached_tensor = None |
| 177 | |
| 178 | |
| 179 | def forward(self, input: torch.Tensor, use_cache = True) -> torch.Tensor: |
| 180 | if use_cache: |
| 181 | if self.cached_tensor is None: |
| 182 | self.cached_tensor = torch.concat([input[:, :, :1]] * 2, dim=2) |
| 183 | input = torch.concat([self.cached_tensor, input], dim=2) |
| 184 | self.cached_tensor = input[:, :, -2:] |
| 185 | return super().forward(input) |
| 186 | |
| 187 | |
| 188 | |