| 46 | """ |
| 47 | |
| 48 | def __init__( |
| 49 | self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, gain=2 |
| 50 | ): |
| 51 | super(WSConv2d, self).__init__() |
| 52 | self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding) |
| 53 | self.scale = (gain / (in_channels * (kernel_size ** 2))) ** 0.5 |
| 54 | self.bias = self.conv.bias |
| 55 | self.conv.bias = None |
| 56 | |
| 57 | # initialize conv layer |
| 58 | nn.init.normal_(self.conv.weight) |
| 59 | nn.init.zeros_(self.bias) |
| 60 | |
| 61 | def forward(self, x): |
| 62 | return self.conv(x * self.scale) + self.bias.view(1, self.bias.shape[0], 1, 1) |