replaces forward method of the original Linear, instead of replacing the original Linear module.
| 20 | |
| 21 | |
| 22 | class LoRAModule(torch.nn.Module): |
| 23 | """ |
| 24 | replaces forward method of the original Linear, instead of replacing the original Linear module. |
| 25 | """ |
| 26 | |
| 27 | def __init__( |
| 28 | self, |
| 29 | lora_name, |
| 30 | org_module: torch.nn.Module, |
| 31 | multiplier=1.0, |
| 32 | lora_dim=4, |
| 33 | alpha=1, |
| 34 | dropout=None, |
| 35 | rank_dropout=None, |
| 36 | module_dropout=None, |
| 37 | ): |
| 38 | """if alpha == 0 or None, alpha is rank (no scaling).""" |
| 39 | super().__init__() |
| 40 | self.lora_name = lora_name |
| 41 | |
| 42 | if org_module.__class__.__name__ == "Conv2d": |
| 43 | in_dim = org_module.in_channels |
| 44 | out_dim = org_module.out_channels |
| 45 | else: |
| 46 | in_dim = org_module.in_features |
| 47 | out_dim = org_module.out_features |
| 48 | |
| 49 | self.lora_dim = lora_dim |
| 50 | if org_module.__class__.__name__ == "Conv2d": |
| 51 | kernel_size = org_module.kernel_size |
| 52 | stride = org_module.stride |
| 53 | padding = org_module.padding |
| 54 | self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False) |
| 55 | self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False) |
| 56 | else: |
| 57 | self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False) |
| 58 | self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False) |
| 59 | |
| 60 | if type(alpha) == torch.Tensor: |
| 61 | alpha = alpha.detach().float().numpy() # without casting, bf16 causes error |
| 62 | alpha = self.lora_dim if alpha is None or alpha == 0 else alpha |
| 63 | self.scale = alpha / self.lora_dim |
| 64 | self.register_buffer("alpha", torch.tensor(alpha)) |
| 65 | |
| 66 | # same as microsoft's |
| 67 | torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5)) |
| 68 | torch.nn.init.zeros_(self.lora_up.weight) |
| 69 | |
| 70 | self.multiplier = multiplier |
| 71 | self.org_module = org_module # remove in applying |
| 72 | self.dropout = dropout |
| 73 | self.rank_dropout = rank_dropout |
| 74 | self.module_dropout = module_dropout |
| 75 | |
| 76 | def apply_to(self): |
| 77 | self.org_forward = self.org_module.forward |
| 78 | self.org_module.forward = self.forward |
| 79 | del self.org_module |
nothing calls this directly
no outgoing calls
no test coverage detected