The layer implements a matrix A and its transpose A^T.
| 17 | |
| 18 | |
| 19 | class Subspace(nn.Module): |
| 20 | """ |
| 21 | The layer implements a matrix A and its transpose A^T. |
| 22 | """ |
| 23 | |
| 24 | def __init__( |
| 25 | self, |
| 26 | in_features: int, |
| 27 | out_features: int, |
| 28 | normalize: bool = False, |
| 29 | orthogonalize: bool = False, |
| 30 | init_norm: float = 1.0, |
| 31 | ): |
| 32 | """ |
| 33 | Args: |
| 34 | in_features (int): |
| 35 | input feature dimension (number of columns) |
| 36 | out_features (int): |
| 37 | output feature dimension (number of rows) |
| 38 | normalize (bool): |
| 39 | whether to normalize the basis vectors (columns of weights). |
| 40 | orthogonalize (bool): |
| 41 | whether to make orthogonal the basis vectors |
| 42 | (columns of weights) by qr decomposition. |
| 43 | Note that it is recommended to turn off (set to False). |
| 44 | init_norm (float): |
| 45 | the l2 norm of the columns |
| 46 | |
| 47 | Notes: |
| 48 | If orthogonalize is `True`, the weight will orthogonalized by QR decomposition, |
| 49 | and thus the shape will change (min of in_feature and out_feature). |
| 50 | """ |
| 51 | super().__init__() |
| 52 | |
| 53 | self.eps = 1e-8 |
| 54 | self.in_features = in_features |
| 55 | self.out_features = out_features |
| 56 | self.normalize = normalize |
| 57 | self.orthogonalize = orthogonalize |
| 58 | self.scale = 1 / math.sqrt(in_features) |
| 59 | |
| 60 | self.weight = nn.Parameter(torch.randn(out_features, in_features)) |
| 61 | |
| 62 | # initialize weights |
| 63 | nn.init.orthogonal_(self.weight, gain=init_norm) |
| 64 | |
| 65 | def __repr__(self): |
| 66 | return ( |
| 67 | f"Subspace({self.in_features}, {self.out_features}, " |
| 68 | f"normalize={self.normalize}, orthogonalize={self.orthogonalize})" |
| 69 | ) |
| 70 | |
| 71 | def forward( |
| 72 | self, |
| 73 | input: torch.Tensor, |
| 74 | mode: str = "A", |
| 75 | input_left: torch.Tensor = None, |
| 76 | ): |