| 51 | self.batch = None |
| 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) |
| 78 | a_out = self.A(x) |
| 79 | b_out = self.B(a_out) |
| 80 | return w_out + b_out * self.lora_scale |
| 81 | |
| 82 | |
| 83 | class LoRAConv(nn.Module): |