| 81 | |
| 82 | |
| 83 | class EqualConv2d(nn.Module): |
| 84 | def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True): |
| 85 | super().__init__() |
| 86 | |
| 87 | self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size)) |
| 88 | self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2) |
| 89 | |
| 90 | self.stride = stride |
| 91 | self.padding = padding |
| 92 | |
| 93 | if bias: |
| 94 | self.bias = nn.Parameter(torch.zeros(out_channel)) |
| 95 | else: |
| 96 | self.bias = None |
| 97 | |
| 98 | def forward(self, input): |
| 99 | |
| 100 | return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding) |
| 101 | |
| 102 | def __repr__(self): |
| 103 | return ( |
| 104 | f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},' |
| 105 | f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})' |
| 106 | ) |
| 107 | |
| 108 | |
| 109 | class EqualLinear(nn.Module): |