File size: 1,417 Bytes
c1a41d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import torch
from tqdm import tqdm
import math
def outlier_channel_split(W, H, mu, to_size):
old_dim = W.shape[-1]
remaining = to_size - old_dim
W = torch.cat([W, torch.zeros(W.shape[0], remaining).to(W.device)], dim=-1)
new_H = torch.zeros(to_size, to_size).to(H.device)
new_H[0:H.shape[0], 0:H.shape[1]] = H
H = new_H
mu = torch.cat([mu, torch.zeros(remaining).to(mu.device)], dim=0)
print('old drange', torch.max(W.flatten()) - torch.min(W.flatten()))
extra_inds = []
dupe_inds = list(range(old_dim))
for i in tqdm(range(old_dim, to_size), desc='outlier channel splitting'):
col = torch.argmax(W.abs()).item() % W.shape[-1]
row = math.ceil(torch.argmax(W.abs()).item() // W.shape[-1])
assert torch.allclose(W[row, col].abs(), torch.max(W.abs().flatten()))
extra_inds.append(col)
dupe_inds.append(dupe_inds[col])
W[:, col] /= 2
W[:, i] = W[:, col]
H[i, 0:i] = H[col, 0:i]
H[0:i, i] = H[0:i, col]
H[i, i] = H[col, col]
mu[i] = mu[col]
i += 1
print('new drange', torch.max(W.flatten()) - torch.min(W.flatten()))
assert torch.allclose(H.cpu(), H.cpu().T)
return W, H, mu, extra_inds, dupe_inds
def fuse_W(W, extra_inds):
for i in range(len(extra_inds)):
W[:, extra_inds[-(i + 1)]] += W[:, -(i + 1)]
return W[:, :W.shape[-1] - len(extra_inds)]
|