| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .general import factorization, FUNC_LIST |
|
|
|
|
| def get_r(oft_blocks, I=None, constraint=0): |
| if I is None: |
| I = torch.eye(oft_blocks.shape[-1], device=oft_blocks.device) |
| if I.ndim < oft_blocks.ndim: |
| for _ in range(oft_blocks.ndim - I.ndim): |
| I = I.unsqueeze(0) |
| |
| q = oft_blocks - oft_blocks.transpose(-1, -2) |
| normed_q = q |
| if constraint is not None and constraint > 0: |
| q_norm = torch.norm(q) + 1e-8 |
| if q_norm > constraint: |
| normed_q = q * constraint / q_norm |
| |
| r = (I + normed_q) @ (I - normed_q).float().inverse() |
| return r |
|
|
|
|
| def weight_gen(org_weight, max_block_size=-1, rescale=False): |
| """### weight_gen |
| |
| Args: |
| org_weight (torch.Tensor): the weight tensor |
| max_block_size (int): max block size |
| rescale (bool, optional): whether to rescale the weight. Defaults to False. |
| |
| Returns: |
| torch.Tensor: oft_blocks[, rescale_weight] |
| """ |
| out_dim, *rest = org_weight.shape |
| block_size, block_num = factorization(out_dim, max_block_size) |
| oft_blocks = torch.zeros(block_num, block_size, block_size) |
| if rescale: |
| return oft_blocks, torch.ones(out_dim, *[1] * len(rest)) |
| else: |
| return oft_blocks, None |
|
|
|
|
| def diff_weight(org_weight, *weights, constraint=None): |
| """### diff_weight |
| |
| Args: |
| org_weight (torch.Tensor): the weight tensor of original model |
| weights (tuple[torch.Tensor]): (oft_blocks[, rescale_weight]) |
| constraint (float, optional): constraint for oft |
| |
| Returns: |
| torch.Tensor: ΔW |
| """ |
| oft_blocks, rescale = weights |
| I = torch.eye(oft_blocks.shape[1], device=oft_blocks.device) |
| r = get_r(oft_blocks, I, constraint) |
|
|
| block_num, block_size, _ = oft_blocks.shape |
| _, *shape = org_weight.shape |
| org_weight = org_weight.to(dtype=r.dtype) |
| org_weight = org_weight.view(block_num, block_size, *shape) |
| |
| weight = torch.einsum( |
| "k n m, k n ... -> k m ...", |
| r - I, |
| org_weight, |
| ).view(-1, *shape) |
| if rescale is not None: |
| weight = rescale * weight |
| weight = weight + (rescale - 1) * org_weight |
| return weight |
|
|
|
|
| def bypass_forward_diff(x, org_out, *weights, constraint=None, need_transpose=False): |
| """### bypass_forward_diff |
| |
| Args: |
| x (torch.Tensor): the input tensor for original model |
| org_out (torch.Tensor): the output tensor from original model |
| weights (tuple[torch.Tensor]): (oft_blocks[, rescale_weight]) |
| constraint (float, optional): constraint for oft |
| need_transpose (bool, optional): |
| whether to transpose the input and output, |
| set to `True` if the original model have "dim" not in the last axis. |
| For example: Convolution layers |
| |
| Returns: |
| torch.Tensor: output tensor |
| """ |
| oft_blocks, rescale = weights |
| block_num, block_size, _ = oft_blocks.shape |
| I = torch.eye(block_size, device=oft_blocks.device) |
| r = get_r(oft_blocks, I, constraint) |
| if need_transpose: |
| org_out = org_out.transpose(1, -1) |
| org_out = org_out.to(dtype=r.dtype) |
| *shape, _ = org_out.shape |
| oft_out = torch.einsum( |
| "k n m, ... k n -> ... k m", r - I, org_out.view(*shape, block_num, block_size) |
| ) |
| out = oft_out.view(*shape, -1) |
| if rescale is not None: |
| out = rescale.transpose(-1, 0) * out |
| out = out + (rescale - 1).transpose(-1, 0) * org_out |
| if need_transpose: |
| out = out.transpose(1, -1) |
| return out |
|
|
