| import torch | |
| import torch.nn as nn | |
| from einops import rearrange | |
| class BayarConv2d(nn.Module): | |
| def __init__( | |
| self, | |
| in_channles: int, | |
| out_channels: int, | |
| kernel_size: int = 5, | |
| stride: int = 1, | |
| padding: int = 0, | |
| magnitude: float = 1.0, | |
| ): | |
| super().__init__() | |
| assert kernel_size > 1, "Bayar conv kernel size must be greater than 1" | |
| self.in_channels = in_channles | |
| self.out_channels = out_channels | |
| self.kernel_size = kernel_size | |
| self.stride = stride | |
| self.padding = padding | |
| self.magnitude = magnitude | |
| self.center_weight = nn.Parameter( | |
| torch.ones(self.in_channels, self.out_channels, 1) * -1.0 * magnitude, | |
| requires_grad=False, | |
| ) | |
| self.kernel_weight = nn.Parameter( | |
| torch.rand((self.in_channels, self.out_channels, kernel_size**2 - 1)), | |
| requires_grad=True, | |
| ) | |
| def _constraint_weight(self): | |
| self.kernel_weight.data = self.kernel_weight.permute(2, 0, 1) | |
| self.kernel_weight.data = torch.div( | |
| self.kernel_weight.data, self.kernel_weight.data.sum(0) | |
| ) | |
| self.kernel_weight.data = self.kernel_weight.permute(1, 2, 0) * self.magnitude | |
| center_idx = self.kernel_size**2 // 2 | |
| full_kernel = torch.cat( | |
| [ | |
| self.kernel_weight[:, :, :center_idx], | |
| self.center_weight, | |
| self.kernel_weight[:, :, center_idx:], | |
| ], | |
| dim=2, | |
| ) | |
| full_kernel = rearrange( | |
| full_kernel, "ci co (kw kh) -> ci co kw kh", kw=self.kernel_size | |
| ) | |
| return full_kernel | |
| def forward(self, x): | |
| x = nn.functional.conv2d( | |
| x, self._constraint_weight(), stride=self.stride, padding=self.padding | |
| ) | |
| return x | |
| if __name__ == "__main__": | |
| device = "cuda" | |
| bayer_conv2d = BayarConv2d(3, 3, 3, magnitude=1).to(device) | |
| bayer_conv2d._constraint_weight() | |
| i = torch.rand(16, 3, 16, 16).to(device) | |
| o = bayer_conv2d(i) | |