File size: 1,294 Bytes
d65ae7d | 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 | from __future__ import annotations
import torch
import torch.nn.functional as F
def centered_classifier_basis(W: torch.Tensor, tol: float = 1e-5, foreground_only: bool = False):
# W: [C,d]
if foreground_only and W.shape[0] > 1:
W = W[1:]
Wc = W - W.mean(dim=0, keepdim=True)
# U basis in feature dimension for row space of Wc
U, S, Vh = torch.linalg.svd(Wc.T, full_matrices=False)
if S.numel() == 0:
return torch.empty(W.shape[1], 0, device=W.device, dtype=W.dtype)
r = int((S > tol * S.max()).sum().item())
return U[:, :max(1, r)]
def random_basis(feature_dim: int, rank: int, device, dtype):
A = torch.randn(feature_dim, rank, device=device, dtype=dtype)
Q, _ = torch.linalg.qr(A, mode="reduced")
return Q
def project_task_and_residual(Fmap: torch.Tensor, Q: torch.Tensor):
# Fmap [B,d,H,W,D], Q [d,r]
if Q.numel() == 0:
return torch.zeros_like(Fmap), Fmap
B, d, H, W, D = Fmap.shape
flat = Fmap.permute(0,2,3,4,1).reshape(-1, d)
task = (flat @ Q) @ Q.T
task = task.reshape(B,H,W,D,d).permute(0,4,1,2,3).contiguous()
residual = Fmap - task
return task, residual
def project_to_residual(x: torch.Tensor, Q: torch.Tensor):
task, res = project_task_and_residual(x, Q)
return res
|