86 lines
2.4 KiB
Python
86 lines
2.4 KiB
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from .general import rebuild_tucker, FUNC_LIST
|
|
|
|
|
|
def weight_gen(org_weight, rank, tucker=True):
|
|
"""### weight_gen
|
|
|
|
Args:
|
|
org_weight (torch.Tensor): the weight tensor
|
|
rank (int): low rank
|
|
|
|
Returns:
|
|
torch.Tensor: down, up[, mid]
|
|
"""
|
|
out_dim, in_dim, *k = org_weight.shape
|
|
if k and tucker:
|
|
down = torch.empty(rank, in_dim, *(1 for _ in k))
|
|
up = torch.empty(out_dim, rank, *(1 for _ in k))
|
|
mid = torch.empty(rank, rank, *k)
|
|
nn.init.kaiming_uniform_(down, a=math.sqrt(5))
|
|
nn.init.constant_(up, 0)
|
|
nn.init.kaiming_uniform_(mid, a=math.sqrt(5))
|
|
return down, up, mid
|
|
else:
|
|
down = torch.empty(rank, in_dim)
|
|
up = torch.empty(out_dim, rank)
|
|
nn.init.kaiming_uniform_(down, a=math.sqrt(5))
|
|
nn.init.constant_(up, 0)
|
|
return down, up, None
|
|
|
|
|
|
def diff_weight(*weights: tuple[torch.Tensor], gamma=1.0):
|
|
"""### diff_weight
|
|
|
|
Get ΔW = BA, where BA is low rank decomposition
|
|
|
|
Args:
|
|
weights (tuple[torch.Tensor]): (down, up[, mid])
|
|
gamma (float, optional): scale factor, normally alpha/rank here
|
|
|
|
Returns:
|
|
torch.Tensor: ΔW
|
|
"""
|
|
d, u, m = weights
|
|
R, I, *k = d.shape
|
|
O, R, *_ = u.shape
|
|
u = u * gamma
|
|
|
|
if m is None:
|
|
result = u.reshape(-1, u.size(1)) @ d.reshape(d.size(0), -1)
|
|
else:
|
|
R, R, *k = m.shape
|
|
u = u.reshape(u.size(0), -1).transpose(0, 1)
|
|
d = d.reshape(d.size(0), -1)
|
|
result = rebuild_tucker(m, u, d)
|
|
return result.reshape(O, I, *k)
|
|
|
|
|
|
def bypass_forward_diff(x, org_out, *weights, gamma=1.0, extra_args={}):
|
|
"""### bypass_forward_diff
|
|
|
|
Args:
|
|
x (torch.Tensor): input tensor
|
|
weights (tuple[torch.Tensor]): (down, up[, mid])
|
|
gamma (float, optional): scale factor, normally alpha/rank here
|
|
extra_args (dict, optional): extra args for forward func, \
|
|
e.g. padding, stride for Conv1/2/3d
|
|
|
|
Returns:
|
|
torch.Tensor: output tensor
|
|
"""
|
|
d, u, m = weights
|
|
if m is not None:
|
|
down = FUNC_LIST[d.dim()](x, d)
|
|
mid = FUNC_LIST[d.dim()](down, m, **extra_args)
|
|
up = FUNC_LIST[d.dim()](mid, u)
|
|
else:
|
|
down = FUNC_LIST[d.dim()](x, d, **extra_args)
|
|
up = FUNC_LIST[d.dim()](down, u)
|
|
return up * gamma
|