| 149 | |
| 150 | |
| 151 | class EqualLinear(nn.Module): |
| 152 | def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None): |
| 153 | super().__init__() |
| 154 | |
| 155 | self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul)) |
| 156 | |
| 157 | if bias: |
| 158 | self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init)) |
| 159 | else: |
| 160 | self.bias = None |
| 161 | |
| 162 | self.activation = activation |
| 163 | |
| 164 | self.scale = (1 / math.sqrt(in_dim)) * lr_mul |
| 165 | self.lr_mul = lr_mul |
| 166 | |
| 167 | def forward(self, input): |
| 168 | |
| 169 | if self.activation: |
| 170 | out = F.linear(input, self.weight * self.scale) |
| 171 | out = fused_leaky_relu(out, self.bias * self.lr_mul) |
| 172 | else: |
| 173 | out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul) |
| 174 | |
| 175 | return out |
| 176 | |
| 177 | def __repr__(self): |
| 178 | return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})') |
| 179 | |
| 180 | |
| 181 | class ScaledLeakyReLU(nn.Module): |