| | import torch |
| | import torch.nn as nn |
| |
|
| | from .base import LycorisBaseModule |
| |
|
| |
|
| | class IA3Module(LycorisBaseModule): |
| | name = "ia3" |
| | support_module = { |
| | "linear", |
| | "conv1d", |
| | "conv2d", |
| | "conv3d", |
| | } |
| | weight_list = ["weight", "on_input"] |
| | weight_list_det = ["on_input"] |
| |
|
| | def __init__( |
| | self, |
| | lora_name, |
| | org_module: nn.Module, |
| | multiplier=1.0, |
| | lora_dim=4, |
| | alpha=1, |
| | dropout=0.0, |
| | rank_dropout=0.0, |
| | module_dropout=0.0, |
| | use_tucker=False, |
| | use_scalar=False, |
| | rank_dropout_scale=False, |
| | weight_decompose=False, |
| | bypass_mode=None, |
| | rs_lora=False, |
| | train_on_input=False, |
| | **kwargs, |
| | ): |
| | """if alpha == 0 or None, alpha is rank (no scaling).""" |
| | super().__init__( |
| | lora_name, |
| | org_module, |
| | multiplier, |
| | dropout, |
| | rank_dropout, |
| | module_dropout, |
| | rank_dropout_scale, |
| | bypass_mode, |
| | ) |
| | if self.module_type not in self.support_module: |
| | raise ValueError(f"{self.module_type} is not supported in IA^3 algo.") |
| |
|
| | if self.module_type.startswith("conv"): |
| | self.isconv = True |
| | in_dim = org_module.in_channels |
| | out_dim = org_module.out_channels |
| | if train_on_input: |
| | train_dim = in_dim |
| | else: |
| | train_dim = out_dim |
| | self.weight = nn.Parameter( |
| | torch.empty(1, train_dim, *(1 for _ in self.shape[2:])) |
| | ) |
| | else: |
| | in_dim = org_module.in_features |
| | out_dim = org_module.out_features |
| | if train_on_input: |
| | train_dim = in_dim |
| | else: |
| | train_dim = out_dim |
| |
|
| | self.weight = nn.Parameter(torch.empty(train_dim)) |
| |
|
| | |
| | torch.nn.init.constant_(self.weight, 0) |
| | self.train_input = train_on_input |
| | self.register_buffer("on_input", torch.tensor(int(train_on_input))) |
| |
|
| | @classmethod |
| | def make_module_from_state_dict(cls, lora_name, orig_module, weight): |
| | module = cls( |
| | lora_name, |
| | orig_module, |
| | 1, |
| | ) |
| | module.weight.data.copy_(weight) |
| | return module |
| |
|
| | def apply_to(self): |
| | self.org_forward = self.org_module[0].forward |
| | self.org_module[0].forward = self.forward |
| |
|
| | def make_weight(self, multiplier=1, shape=None, device=None, diff=False): |
| | weight = self.weight * multiplier + int(not diff) |
| | if self.train_input: |
| | diff = self.org_weight * weight |
| | else: |
| | diff = self.org_weight.transpose(0, 1) * weight |
| | diff = diff.transpose(0, 1) |
| | if shape is not None: |
| | diff = diff.view(shape) |
| | if device is not None: |
| | diff = diff.to(device) |
| | return diff |
| |
|
| | def get_diff_weight(self, multiplier=1, shape=None, device=None): |
| | diff = self.make_weight( |
| | multiplier=multiplier, shape=shape, device=device, diff=True |
| | ) |
| | return diff, None |
| |
|
| | def get_merged_weight(self, multiplier=1, shape=None, device=None): |
| | diff = self.make_weight(multiplier=multiplier, shape=shape, device=device) |
| | return diff, None |
| |
|
| | def _bypass_forward(self, x, scale=1, diff=False): |
| | weight = self.weight * scale + int(not diff) |
| | if self.train_input: |
| | x = x * weight |
| | out = self.org_forward(x) |
| | if not self.train_input: |
| | out = out * weight |
| | return out |
| |
|
| | def bypass_forward_diff(self, x, scale=1): |
| | return self._bypass_forward(x, scale, diff=True) |
| |
|
| | def bypass_forward(self, x, scale=1): |
| | return self._bypass_forward(x, scale, diff=False) |
| |
|
| | def forward(self, x, *args, **kwargs): |
| | if self.module_dropout and self.training: |
| | if torch.rand(1) < self.module_dropout: |
| | return self.org_forward(x) |
| | if self.bypass_mode: |
| | return self.bypass_forward(x, self.multiplier) |
| | else: |
| | weight = self.get_merged_weight(multiplier=self.multiplier)[0] |
| | bias = ( |
| | None |
| | if self.org_module[0].bias is None |
| | else self.org_module[0].bias.data |
| | ) |
| | return self.op(x, weight, bias, **self.kw_dict) |
| |
|
