ncut-pytorch 3.0.0.dev4__tar.gz → 3.0.0.dev6__tar.gz

{ncut_pytorch-3.0.0.dev4 → ncut_pytorch-3.0.0.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ncut_pytorch
-Version: 3.0.0.dev4
+Version: 3.0.0.dev6
 Summary: Normalized Cut and Spectral Embedding
 Author-email: Huzheng Yang <huze.yann@gmail.com>
 License-Expression: MIT

{ncut_pytorch-3.0.0.dev4 → ncut_pytorch-3.0.0.dev6}/ncut_pytorch/ncuts/ncut_click.py

@@ -89,7 +89,6 @@ def ncut_click_prompt(
         nystrom_X,
         n_neighbors=config.n_neighbors,
         n_sample=config.n_sample2,
-        matmul_chunk_size=config.matmul_chunk_size,
         device=device,
         return_indices=True,
     )

{ncut_pytorch-3.0.0.dev4 → ncut_pytorch-3.0.0.dev6}/ncut_pytorch/ncuts/ncut_nystrom.py

@@ -9,7 +9,10 @@ from ncut_pytorch.utils.math import rbf_affinity, cosine_affinity
 from ncut_pytorch.utils.math import gram_schmidt, normalize_affinity, grad_safe_eig_solve, correct_rotation, keep_topk_per_row, svd_lowrank
 from ncut_pytorch.utils.sample import farthest_point_sampling
 from ncut_pytorch.utils.device import auto_device
+import logging
 
+MATMUL_CHUNK_SIZE = 65536
+SMALL_SCALE_THRESHOLD = 8192    # if the number of nodes is less than SMALL_SCALE_THRESHOLD, skip nystrom approximation use exact ncut
 
 class NystromConfig:
     """
@@ -21,7 +24,6 @@ class NystromConfig:
     n_sample2 = 1024                # number of samples for eigenvector propagation, 1024 is large enough for most cases
     n_neighbors = 32                # number of neighbors for eigenvector propagation, 10 is large enough for most cases
     n_neighbors_max_ratio = 1/32    # max ratio of n_neighbors to n_sample2, to avoid over smoothing
-    matmul_chunk_size = 65536       # chunk size for matrix multiplication, larger chunk size is faster but requires more memory
     
     def update(self, kwargs: dict):
         for key, value in kwargs.items():
@@ -78,9 +80,11 @@ def ncut_fn(
     device = auto_device(X.device, device)
 
     # subsample for nystrom approximation
-    is_enough_data = X.shape[0] > config.n_sample
     n_sample = min(config.n_sample, int(X.shape[0]*config.n_sample_max_ratio))
-    nystrom_indices = farthest_point_sampling(X, n_sample=n_sample, device=device) if is_enough_data else np.arange(X.shape[0])
+    if X.shape[0] > SMALL_SCALE_THRESHOLD:
+        nystrom_indices = farthest_point_sampling(X, n_sample=n_sample, device=device)
+    else:
+        nystrom_indices = torch.arange(X.shape[0])
     nystrom_X = X[nystrom_indices].to(device)
 
     sigma, repulsion_sigma = find_optimal_sigma(nystrom_X, quantile_sigma, quantile_sigma_repulsion, sigma, repulsion_sigma, affinity_fn)
@@ -95,10 +99,6 @@ def ncut_fn(
     if no_propagation:
         return nystrom_eigvec, eigval, nystrom_indices, sigma
 
-    if not is_enough_data:
-        # skip nystrom approximation if not enough data, use exact ncut
-        return nystrom_eigvec, eigval
-
     # propagate eigenvectors from subgraph to full graph
     eigvec = nystrom_propagate(
         nystrom_eigvec,
@@ -107,7 +107,6 @@ def ncut_fn(
         extrapolation_factor=extrapolation_factor,
         n_neighbors=config.n_neighbors,
         n_sample=config.n_sample2,
-        matmul_chunk_size=config.matmul_chunk_size,
         device=device,
     )
 
@@ -117,6 +116,7 @@ def ncut_fn(
 
     return eigvec, eigval
     
+
 def find_optimal_sigma(
     X: torch.Tensor,
     quantile_sigma: float = 0.25,
@@ -137,6 +137,7 @@ def find_optimal_sigma(
             raise ValueError(f"`sigma` need to be provided for affinity function {affinity_fn}, (sigma=0.5, repulsion_sigma=0.3)")
     return sigma, repulsion_sigma
 
+
 def ncut_with_repulsion(
     X: torch.Tensor,
     n_eig: int = 100,
@@ -197,11 +198,16 @@ def nystrom_propagate(
         nystrom_X (torch.Tensor): input features from nystrom sampled nodes, shape (m, D)
         extrapolation_factor (float): control how far can we extrapolate, larger extrapolation_factor means we can extrapolate further, default 1.0
         device (str): device to use for computation, if 'auto', will detect GPU automatically
-        affinity_fn (callable): affinity function, default rbf_affinity. Should accept (X1, X2=None, sigma=float) and return affinity matrix
+        return_indices (bool): whether to return the indices used for propagation
 
     Returns:
         torch.Tensor: output propagated by nearest neighbors, shape (N, D)
     """
+    if X.shape[0] <= SMALL_SCALE_THRESHOLD and nystrom_out.shape == X.shape and torch.allclose(nystrom_X.to(X.device), X, atol=1e-6):
+        # skip propagation if nystrom_out is the same as X, for small scale graph that don't need nystrom approximation
+        if return_indices:
+            return nystrom_out, np.arange(X.shape[0])
+        return nystrom_out
 
     config = NystromConfig()
     config.update(kwargs)
@@ -217,33 +223,36 @@ def nystrom_propagate(
     
     D = rbf_affinity(nystrom_X, sigma=sigma).mean(1)
 
-    all_outs = []
-    n_chunk = config.matmul_chunk_size
     n_neighbors = int(min(config.n_neighbors, len(indices)*config.n_neighbors_max_ratio))
     n_neighbors = max(n_neighbors, 4)
+    n_chunk = _find_max_chunk_size(X, nystrom_X, device)
+
+    all_outs = torch.empty((X.shape[0], nystrom_out.shape[-1]), device=output_device, dtype=nystrom_out.dtype)
     for i in range(0, X.shape[0], n_chunk):
         end = min(i + n_chunk, X.shape[0])
 
         _Ai = rbf_affinity(X[i:end].to(device), nystrom_X, sigma=sigma)
         _Ai, _indices = keep_topk_per_row(_Ai, n_neighbors)  # (n, n_neighbors)
+        
         _Di = D[_indices].sum(1)
         _Ai = _Ai / _Di[:, None]
 
-        weights = _Ai[..., None]  # (n, n_neighbors, 1)
-        neighbors = nystrom_out[_indices.flatten()]
-        neighbors = neighbors.reshape(-1, n_neighbors, nystrom_out.shape[-1])  # (n, n_neighbors, d)
-        out = weights * neighbors  # (n, n_neighbors, d)
-        out = out.sum(dim=1)  # (n, d)
-
-        out = out.to(output_device)
-        all_outs.append(out)
+        out = torch.einsum('nk,nkd->nd', _Ai, nystrom_out[_indices])
 
-    all_outs = torch.cat(all_outs, dim=0)
+        all_outs[i:end] = out.to(output_device)
 
     if return_indices:
         return all_outs, indices
-    
     return all_outs
 
 
-
+def _find_max_chunk_size(X: torch.Tensor, nystrom_X: torch.Tensor, device: str):
+    max_chunk_size = MATMUL_CHUNK_SIZE
+    while max_chunk_size > 1:
+        try:
+            _ = rbf_affinity(X[:max_chunk_size].to(device), nystrom_X)
+            return max_chunk_size
+        except RuntimeError as e:
+            max_chunk_size = max_chunk_size // 2
+            continue
+    raise RuntimeError("failed to find max chunk size")

{ncut_pytorch-3.0.0.dev4 → ncut_pytorch-3.0.0.dev6}/ncut_pytorch/utils/grad.py

@@ -113,3 +113,142 @@ def rbf_eigvec_manual_grad(
 
     return grad_u
 
+
+class MultiSpectralProjectorFromMasks(torch.autograd.Function):
+    """
+    A (symmetric) -> {P_b}_b, where P_b = U_{S_b} U_{S_b}^T and S_b is specified by a boolean mask.
+
+    Computes eigh(A) ONCE, sorts eigenpairs DESCENDING (largest-first), then forms projectors
+    for each mask.
+
+    Inputs:
+      A:     [N,N] (float), symmetric (or will be symmetrized if symmetrize=True)
+      masks: [B,N] (bool), masks[b,i]=True selects eigenvector i in DESCENDING eigen-order.
+
+    Output:
+      P: [B,N,N]
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        A: torch.Tensor,          # [N,N]
+        masks: torch.Tensor,      # [B,N] bool (in DESCENDING eigen-order)
+        gap_eps: float = 0.0,
+        symmetrize: bool = True,
+    ):
+        if A.ndim != 2 or A.shape[0] != A.shape[1]:
+            raise ValueError(f"A must be square [N,N], got {tuple(A.shape)}")
+        if masks.ndim != 2:
+            raise ValueError(f"masks must be [B,N], got {tuple(masks.shape)}")
+        if masks.dtype != torch.bool:
+            raise ValueError("masks must be boolean")
+        N = A.shape[0]
+        B, N2 = masks.shape
+        if N2 != N:
+            raise ValueError(f"masks second dim must equal N={N}, got {N2}")
+        if (masks.sum(dim=1) == 0).any():
+            raise ValueError("Each mask row must select at least one eigenvector.")
+
+        device = A.device
+        masks = masks.to(device=device)
+
+        A_used = 0.5 * (A + A.T) if symmetrize else A
+
+        # eigh ascending -> flip to descending
+        evals_asc, U_asc = torch.linalg.eigh(A_used)
+        evals = torch.flip(evals_asc, dims=[0])   # [N] descending
+        U = torch.flip(U_asc, dims=[1])           # [N,N] descending columns
+
+        # Build projectors
+        P_out = []
+        for b in range(B):
+            U_S = U[:, masks[b]]                  # [N,p_b]
+            P_b = U_S @ U_S.T                     # [N,N]
+            P_out.append(P_b)
+        P = torch.stack(P_out, dim=0)             # [B,N,N]
+
+        ctx.save_for_backward(U, evals, masks)
+        ctx.gap_eps = float(gap_eps)
+        ctx.symmetrize = bool(symmetrize)
+
+        return P
+
+    @staticmethod
+    def backward(ctx, grad_P: torch.Tensor):
+        U, evals, masks = ctx.saved_tensors
+        gap_eps = ctx.gap_eps
+        symmetrize = ctx.symmetrize
+
+        if grad_P.ndim != 3:
+            raise ValueError(f"grad_P must be [B,N,N], got {tuple(grad_P.shape)}")
+        B, N, N2 = grad_P.shape
+        if N != N2:
+            raise ValueError("grad_P must be square per batch")
+
+        grad_A_used = torch.zeros((N, N), device=grad_P.device, dtype=grad_P.dtype)
+
+        for b in range(B):
+            mask = masks[b]                       # [N]
+            U_S = U[:, mask]                      # [N,p]
+            U_perp = U[:, ~mask]                  # [N,N-p]
+            lam_S = evals[mask]                   # [p]
+            lam_perp = evals[~mask]               # [N-p]
+
+            # symmetric part only matters
+            G = grad_P[b]
+            Gs = 0.5 * (G + G.T)
+
+            # H = U_perp^T Gs U_S
+            H = U_perp.T @ (Gs @ U_S)             # [N-p,p]
+
+            denom = lam_S[None, :] - lam_perp[:, None]  # [N-p,p]
+            if gap_eps > 0.0:
+                denom = torch.sign(denom) * torch.clamp(denom.abs(), min=gap_eps)
+
+            Q = H / denom                          # [N-p,p]
+
+            Bmat = U_perp @ (Q @ U_S.T)            # [N,N]
+            grad_A_used = grad_A_used + (Bmat + Bmat.T)
+
+        grad_A = 0.5 * (grad_A_used + grad_A_used.T) if symmetrize else grad_A_used
+        return grad_A, None, None, None
+
+
+def spectral_projectors_from_masks(
+    A: torch.Tensor,
+    masks: torch.Tensor,
+    gap_eps: float = 0.0,
+    symmetrize: bool = True,
+):
+    """
+    Convenience wrapper.
+
+    masks: [B,N] bool in DESCENDING eigen-order (0 = largest eigenvalue).
+    returns P: [B,N,N]
+    """
+    return MultiSpectralProjectorFromMasks.apply(A, masks, gap_eps, symmetrize)
+
+
+if __name__ == "__main__":
+    B = 2
+    N = 1000
+    masks = torch.zeros(B, N, dtype=torch.bool)
+    masks[0, :3] = True          # top-3 eigenvectors (largest-first)
+    masks[1, 3:6] = True         # next-3
+
+
+    A1 = torch.randn(N, N)
+    A1 = 0.5 * (A1 + A1.T)
+    A1.requires_grad_(True)
+    P1 = spectral_projectors_from_masks(A1, masks)
+
+    A2 = torch.randn(N, N)
+    A2 = 0.5 * (A2 + A2.T)
+    A2.requires_grad_(True)
+    P2 = spectral_projectors_from_masks(A2, masks)
+
+    loss = torch.norm(P1 - P2, p=2, dim=(0, 1)).sum()
+    loss.backward()
+    print(A1.grad.shape)    
+    print(A2.grad.shape)

{ncut_pytorch-3.0.0.dev4 → ncut_pytorch-3.0.0.dev6}/ncut_pytorch/utils/math.py

@@ -45,7 +45,11 @@ def rbf_affinity(
     sigma = sigma if gamma is None else check_gamma_deprecated(gamma)
     X2 = X1 if X2 is None else X2
 
-    dist2 = torch.cdist(X1, X2, p=2)**2
+    try:
+        dist2 = torch.cdist(X1, X2, p=2)**2
+    except NotImplementedError:
+        dist2 = X1.unsqueeze(1) - X2.unsqueeze(0)
+        dist2 = dist2.pow(2).sum(dim=-1)
     W = torch.exp(-dist2 / (2.0 * sigma * sigma))   # [N,M]
     if zero_diag and X1 is X2:
         W = W.clone()
@@ -98,11 +102,10 @@ def grad_safe_eig_solve(
         is_symmetric = mat.shape[0] == mat.shape[1]
         if is_symmetric:
             s, u = torch.linalg.eigh(mat)
+            s = torch.flip(s, dims=[0])
+            u = torch.flip(u, dims=[1])
         else:
             s, u = torch.linalg.eig(mat)
-        sort_idx = torch.argsort(s, dim=0, descending=True)
-        s = s[sort_idx]
-        u = u[:, sort_idx]
         return u.to(dtype), s.to(dtype), None
 
     try:

{ncut_pytorch-3.0.0.dev4 → ncut_pytorch-3.0.0.dev6}/ncut_pytorch/utils/sample.py

@@ -60,5 +60,5 @@ def _farthest_point_sampling(
     assert not torch.any(torch.isnan(X)), "X contains NaN"
     assert not torch.any(torch.isinf(X)), "X contains Inf"
     
-    samples_idx = sample_idx(X.cpu(), n_sample).numpy()
+    samples_idx = sample_idx(X.cpu(), n_sample)
     return samples_idx

{ncut_pytorch-3.0.0.dev4 → ncut_pytorch-3.0.0.dev6}/ncut_pytorch.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ncut_pytorch
-Version: 3.0.0.dev4
+Version: 3.0.0.dev6
 Summary: Normalized Cut and Spectral Embedding
 Author-email: Huzheng Yang <huze.yann@gmail.com>
 License-Expression: MIT

{ncut_pytorch-3.0.0.dev4 → ncut_pytorch-3.0.0.dev6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ncut_pytorch"
-version = "3.0.0dev4"
+version = "3.0.0dev6"
 authors = [
     { name = "Huzheng Yang", email = "huze.yann@gmail.com" },
 ]