| 52 | |
| 53 | class LoraLinear(torch.nn.Module): |
| 54 | def __init__( |
| 55 | self, |
| 56 | out_features, |
| 57 | in_features, |
| 58 | rank = None, |
| 59 | lora_scale = 1.0, |
| 60 | ): |
| 61 | super().__init__() |
| 62 | self.rank = rank |
| 63 | self.lora_scale = lora_scale |
| 64 | |
| 65 | # original weight of the matrix |
| 66 | self.W = nn.Linear(in_features, out_features, bias=False) |
| 67 | for p in self.W.parameters(): |
| 68 | p.requires_grad_(False) |
| 69 | |
| 70 | self.A = nn.Linear(in_features, rank, bias=False) |
| 71 | self.B = nn.Linear(rank, out_features, bias=False) |
| 72 | # b should be init wiht 0 |
| 73 | for p in self.B.parameters(): |
| 74 | p.detach().zero_() |
| 75 | |
| 76 | def forward(self, x): |
| 77 | w_out = self.W(x) |