(self, C, kernel_size=3, bias=False)
| 24 | class IdentityConv2d(nn.Conv2d): |
| 25 | """Same-shape Conv2d initialized to identity (Dirac).""" |
| 26 | def __init__(self, C, kernel_size=3, bias=False): |
| 27 | pad = kernel_size // 2 |
| 28 | super().__init__(C, C, kernel_size, padding=pad, bias=bias) |
| 29 | with torch.no_grad(): |
| 30 | init.dirac_(self.weight) |
| 31 | if self.bias is not None: |
| 32 | self.bias.zero_() |
| 33 | |
| 34 | def conv(n_in, n_out, **kwargs): |
| 35 | return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) |