166 lines
5.2 KiB
Python
166 lines
5.2 KiB
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from .general import FUNC_LIST
|
|
|
|
|
|
class HadaWeight(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, w1d, w1u, w2d, w2u, scale=torch.tensor(1)):
|
|
ctx.save_for_backward(w1d, w1u, w2d, w2u, scale)
|
|
diff_weight = ((w1u @ w1d) * (w2u @ w2d)) * scale
|
|
return diff_weight
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_out):
|
|
(w1d, w1u, w2d, w2u, scale) = ctx.saved_tensors
|
|
grad_out = grad_out * scale
|
|
temp = grad_out * (w2u @ w2d)
|
|
grad_w1u = temp @ w1d.T
|
|
grad_w1d = w1u.T @ temp
|
|
|
|
temp = grad_out * (w1u @ w1d)
|
|
grad_w2u = temp @ w2d.T
|
|
grad_w2d = w2u.T @ temp
|
|
|
|
del temp
|
|
return grad_w1d, grad_w1u, grad_w2d, grad_w2u, None
|
|
|
|
|
|
class HadaWeightTucker(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, t1, w1d, w1u, t2, w2d, w2u, scale=torch.tensor(1)):
|
|
ctx.save_for_backward(t1, w1d, w1u, t2, w2d, w2u, scale)
|
|
|
|
rebuild1 = torch.einsum("i j ..., j r, i p -> p r ...", t1, w1d, w1u)
|
|
rebuild2 = torch.einsum("i j ..., j r, i p -> p r ...", t2, w2d, w2u)
|
|
|
|
return rebuild1 * rebuild2 * scale
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_out):
|
|
(t1, w1d, w1u, t2, w2d, w2u, scale) = ctx.saved_tensors
|
|
grad_out = grad_out * scale
|
|
|
|
temp = torch.einsum("i j ..., j r -> i r ...", t2, w2d)
|
|
rebuild = torch.einsum("i j ..., i r -> r j ...", temp, w2u)
|
|
|
|
grad_w = rebuild * grad_out
|
|
del rebuild
|
|
|
|
grad_w1u = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
|
|
grad_temp = torch.einsum("i j ..., i r -> r j ...", grad_w, w1u.T)
|
|
del grad_w, temp
|
|
|
|
grad_w1d = torch.einsum("i r ..., i j ... -> r j", t1, grad_temp)
|
|
grad_t1 = torch.einsum("i j ..., j r -> i r ...", grad_temp, w1d.T)
|
|
del grad_temp
|
|
|
|
temp = torch.einsum("i j ..., j r -> i r ...", t1, w1d)
|
|
rebuild = torch.einsum("i j ..., i r -> r j ...", temp, w1u)
|
|
|
|
grad_w = rebuild * grad_out
|
|
del rebuild
|
|
|
|
grad_w2u = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
|
|
grad_temp = torch.einsum("i j ..., i r -> r j ...", grad_w, w2u.T)
|
|
del grad_w, temp
|
|
|
|
grad_w2d = torch.einsum("i r ..., i j ... -> r j", t2, grad_temp)
|
|
grad_t2 = torch.einsum("i j ..., j r -> i r ...", grad_temp, w2d.T)
|
|
del grad_temp
|
|
return grad_t1, grad_w1d, grad_w1u, grad_t2, grad_w2d, grad_w2u, None
|
|
|
|
|
|
def make_weight(w1d, w1u, w2d, w2u, scale):
|
|
return HadaWeight.apply(w1d, w1u, w2d, w2u, scale)
|
|
|
|
|
|
def make_weight_tucker(t1, w1d, w1u, t2, w2d, w2u, scale):
|
|
return HadaWeightTucker.apply(t1, w1d, w1u, t2, w2d, w2u, scale)
|
|
|
|
|
|
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: w1d, w2d, w1u, w2u[, t1, t2]
|
|
"""
|
|
out_dim, in_dim, *k = org_weight.shape
|
|
if k and tucker:
|
|
w1d = torch.empty(rank, in_dim)
|
|
w1u = torch.empty(rank, out_dim)
|
|
t1 = torch.empty(rank, rank, *k)
|
|
w2d = torch.empty(rank, in_dim)
|
|
w2u = torch.empty(rank, out_dim)
|
|
t2 = torch.empty(rank, rank, *k)
|
|
nn.init.normal_(t1, std=0.1)
|
|
nn.init.normal_(t2, std=0.1)
|
|
else:
|
|
w1d = torch.empty(rank, in_dim)
|
|
w1u = torch.empty(out_dim, rank)
|
|
w2d = torch.empty(rank, in_dim)
|
|
w2u = torch.empty(out_dim, rank)
|
|
t1 = t2 = None
|
|
nn.init.normal_(w1d, std=1)
|
|
nn.init.constant_(w1u, 0)
|
|
nn.init.normal_(w2d, std=1)
|
|
nn.init.normal_(w2u, std=0.1)
|
|
return w1d, w1u, w2d, w2u, t1, t2
|
|
|
|
|
|
def diff_weight(*weights, gamma=1.0):
|
|
"""### diff_weight
|
|
|
|
Get ΔW = BA, where BA is low rank decomposition
|
|
|
|
Args:
|
|
wegihts (tuple[torch.Tensor]): (w1d, w2d, w1u, w2u[, t1, t2])
|
|
gamma (float, optional): scale factor, normally alpha/rank here
|
|
|
|
Returns:
|
|
torch.Tensor: ΔW
|
|
"""
|
|
w1d, w1u, w2d, w2u, t1, t2 = weights
|
|
if t1 is not None and t2 is not None:
|
|
R, I = w1d.shape
|
|
R, O = w1u.shape
|
|
R, R, *k = t1.shape
|
|
result = make_weight_tucker(t1, w1d, w1u, t2, w2d, w2u, gamma)
|
|
else:
|
|
R, I, *k = w1d.shape
|
|
O, R, *_ = w1u.shape
|
|
w1d = w1d.reshape(w1d.size(0), -1)
|
|
w1u = w1u.reshape(-1, w1u.size(1))
|
|
w2d = w2d.reshape(w2d.size(0), -1)
|
|
w2u = w2u.reshape(-1, w2u.size(1))
|
|
result = make_weight(w1d, w1u, w2d, w2u, gamma)
|
|
|
|
result = result.reshape(O, I, *k)
|
|
return result
|
|
|
|
|
|
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]): (w1d, w2d, w1u, w2u[, t1, t2])
|
|
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
|
|
"""
|
|
w1d, w1u, w2d, w2u, t1, t2 = weights
|
|
diff_w = diff_weight(w1d, w1u, w2d, w2u, t1, t2, gamma)
|
|
return FUNC_LIST[w1d.dim() if t1 is None else t1.dim()](x, diff_w, **extra_args)
|