| 123 | |
| 124 | |
| 125 | class EqualConv2d(nn.Module): |
| 126 | def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True): |
| 127 | super().__init__() |
| 128 | |
| 129 | self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size)) |
| 130 | self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2) |
| 131 | |
| 132 | self.stride = stride |
| 133 | self.padding = padding |
| 134 | |
| 135 | if bias: |
| 136 | self.bias = nn.Parameter(torch.zeros(out_channel)) |
| 137 | else: |
| 138 | self.bias = None |
| 139 | |
| 140 | def forward(self, input): |
| 141 | |
| 142 | return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding, ) |
| 143 | |
| 144 | def __repr__(self): |
| 145 | return ( |
| 146 | f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},' |
| 147 | f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})' |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | class EqualLinear(nn.Module): |