r""" The layer implements a typical linear layer but allows the input x to be shifted by dx. .. math:: & y = W * (x + dx) + b + b0, \text{ where} \\ & \quad * \text{ is matrix-vector multiplication} \\ & \quad x \in R^{cin}, \\ & \quad dx \in R^{cin}, \\
| 301 | |
| 302 | |
| 303 | class ShiftedLinearLayer(nn.Module): |
| 304 | r""" |
| 305 | The layer implements a typical linear layer |
| 306 | but allows the input x to be shifted by dx. |
| 307 | |
| 308 | .. math:: |
| 309 | & y = W * (x + dx) + b + b0, \text{ where} \\ |
| 310 | & \quad * \text{ is matrix-vector multiplication} \\ |
| 311 | & \quad x \in R^{cin}, \\ |
| 312 | & \quad dx \in R^{cin}, \\ |
| 313 | & \quad W \in R^{cout \times cin}, |
| 314 | """ |
| 315 | |
| 316 | def __init__( |
| 317 | self, |
| 318 | in_features: int, |
| 319 | out_features: int, |
| 320 | bias=True, |
| 321 | fixed_bias: float = None, |
| 322 | bias_init_val: float = 0.0, |
| 323 | lr_multiplier: float = 1.0, |
| 324 | demodulate=False, |
| 325 | ): |
| 326 | r""" |
| 327 | Args: |
| 328 | in_features (int): |
| 329 | input feature dimension |
| 330 | out_features (int): |
| 331 | output feature dimension |
| 332 | bias (bool): |
| 333 | whether to learn bias :math:`b` |
| 334 | fixed_bias (float): |
| 335 | a fixed bias b0 added after Wx + b + b0. |
| 336 | lr_multiplier (float): |
| 337 | a factor controls the learning rate of the layer. |
| 338 | demodulate (bool): |
| 339 | whether to normalize the row of W. |
| 340 | """ |
| 341 | super().__init__() |
| 342 | |
| 343 | self.eps = 1e-8 |
| 344 | self.in_features = in_features |
| 345 | self.out_features = out_features |
| 346 | self.fixed_bias = fixed_bias |
| 347 | self.demodulate = demodulate |
| 348 | self.lr_multiplier = lr_multiplier |
| 349 | self.scale = 1 / math.sqrt(in_features) |
| 350 | |
| 351 | self.weight = nn.Parameter(torch.randn(out_features, in_features)) |
| 352 | if bias: |
| 353 | self.bias = nn.Parameter(torch.ones(out_features) * bias_init_val) |
| 354 | else: |
| 355 | self.bias = None |
| 356 | |
| 357 | def __repr__(self): |
| 358 | return ( |
| 359 | f"ShiftedLinearLayer({self.in_features}, {self.out_features}, " |
| 360 | f"bias={self.bias is not None}, fixed_bias={self.fixed_bias})" |