ncut-pytorch 3.0.0.dev0__tar.gz → 3.0.0.dev2__tar.gz

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/PKG-INFO

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

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/ncut_pytorch/color/coloring.py

@@ -218,7 +218,6 @@ def _nystrom_dimension_reduction(
         X[subgraph_indices],
         n_neighbors=knn,
         device=device,
-        move_output_to_cpu=True,
     ))
     rgb = rgb_func(X_nd, q)
     return X_nd, rgb

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/ncut_pytorch/color/mspace.py

@@ -56,7 +56,7 @@ def ncut_wrapper(features, n_eig, sigma=None):
 
     # features.requires_grad_(True)
     sigma = sigma or features.std(0).sum().item()
-    # eigvec, eigval = ncut_fn(features, n_eig, sigma=sigma, track_grad=True)
+    eigvec, eigval = ncut_fn(features, n_eig, sigma=sigma)
     W = rbf_affinity(features, sigma=sigma)
     # W = cosine_affinity(features, sigma=1.0)
     A = normalize_affinity(W)

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/ncut_pytorch/ncut.py

@@ -14,28 +14,30 @@ class Ncut:
     def __init__(
             self,
             n_eig: int = 100,
-            track_grad: bool = False,
-            d_sigma: float = None,
-            sigma: float = None,
-            repulsion_sigma: float = None,
-            repulsion_weight: float = 0.2,
-            extrapolation_factor: float = 1.0,
-            device: str = None,
+            quantile_sigma: float = 0.25,
+            quantile_sigma_repulsion: float = 0.20,
+            sigma: float | None = None,
+            repulsion_sigma: float | None = None,
+            repulsion_weight: float | None = None,
             affinity_fn: Union["rbf_affinity", "cosine_affinity"] = rbf_affinity,
+            extrapolation_factor: float = 1.0,
+            exact_gradient: bool = False,
+            device: str | None = None,
             **kwargs,
     ):
         """
         
         Args:       
             n_eig (int): number of eigenvectors
-            track_grad (bool): keep track of pytorch gradients
-            d_sigma (float): affinity sigma parameter, lower d_sigma results in a sharper eigenvectors
+            n_eig (int): number of eigenvectors
+            quantile_sigma (float): quantile of affinity sigma parameter, lower quantile_sigma results in sharper eigenvectors
+            quantile_sigma_repulsion (float): quantile of repulsion sigma parameter, lower quantile_sigma_repulsion results in sharper eigenvectors
             sigma (float): affinity parameter, override d_sigma if provided
             repulsion_sigma (float): (if use repulsion) repulsion sigma parameter, default None (no repulsion)
             repulsion_weight (float): (if use repulsion) repulsion weight, default 0.2
-            extrapolation_factor (float): control how far can we extrapolate, larger extrapolation_factor means we can extrapolate further, default 1.0
-            device (str): device, default 'auto' (auto detect GPU)
             affinity_fn (callable): affinity function, default rbf_affinity. Should accept (X1, X2=None, sigma=float) and return affinity matrix
+            extrapolation_factor (float): control how far can we extrapolate, larger extrapolation_factor means we can extrapolate further, default 1.0
+            exact_gradient (bool): use full spectrum and exact gradient, can be slower and unstable, default False            device (str): device, default 'auto' (auto detect GPU)
             
         Examples:
             >>> from ncut_pytorch import Ncut
@@ -52,13 +54,14 @@ class Ncut:
             >>> print(new_eigvec.shape)  # (500, 20)
         """
         self.n_eig = n_eig
-        self.d_sigma = d_sigma
+        self.quantile_sigma = quantile_sigma
+        self.quantile_sigma_repulsion = quantile_sigma_repulsion
         self.sigma = sigma
         self.repulsion_sigma = repulsion_sigma
         self.repulsion_weight = repulsion_weight
         self.extrapolation_factor = extrapolation_factor
+        self.exact_gradient = exact_gradient
         self.device = device
-        self.track_grad = track_grad
         self.affinity_fn = affinity_fn
         self.kwargs = kwargs
 
@@ -83,12 +86,13 @@ class Ncut:
             ncut_fn(
                 X,
                 n_eig=self.n_eig,
-                d_sigma=self.d_sigma,
+                quantile_sigma=self.quantile_sigma,
+                quantile_sigma_repulsion=self.quantile_sigma_repulsion,
                 sigma=self.sigma,
                 repulsion_sigma=self.repulsion_sigma,
                 repulsion_weight=self.repulsion_weight,
                 device=self.device,
-                track_grad=self.track_grad,
+                exact_gradient=self.exact_gradient,
                 no_propagation=True,
                 affinity_fn=self.affinity_fn,
                 **self.kwargs
@@ -121,7 +125,6 @@ class Ncut:
             self._nystrom_x,
             extrapolation_factor=self.extrapolation_factor,
             device=self.device,
-            track_grad=self.track_grad,
             **self.kwargs
         )
         return eigvec
@@ -137,15 +140,5 @@ class Ncut:
         """
         return self.fit(X).transform(X)
 
-    def __new__(cls, X: torch.Tensor = None, n_eig: int = 100, track_grad: bool = False, d_sigma: float = None,
-                device: str = None, affinity_fn: Callable[[torch.Tensor, torch.Tensor, float], torch.Tensor] = rbf_affinity, 
-                **kwargs) -> Union["Ncut", torch.Tensor]:
-        if X is not None:
-            # function-like behavior
-            eigvec, eigval = ncut_fn(X, n_eig=n_eig, track_grad=track_grad, d_sigma=d_sigma, device=device, affinity_fn=affinity_fn, **kwargs)
-            return eigvec
-        # normal class instantiation
-        return super().__new__(cls)
-
     def __call__(self, X: torch.Tensor) -> torch.Tensor:
         return self.fit_transform(X)

ncut_pytorch-3.0.0.dev2/ncut_pytorch/ncuts/ncut_click.py

@@ -0,0 +1,102 @@
+__all__ = ['ncut_click_prompt']
+
+from typing import Callable, Union
+
+import numpy as np
+import torch
+
+from ncut_pytorch.utils.sigma import find_sigma_by_degree
+from ncut_pytorch.utils.math import rbf_affinity, cosine_affinity, normalize_affinity
+from ncut_pytorch.utils.sample import farthest_point_sampling
+from ncut_pytorch.utils.device import auto_device
+from .ncut_nystrom import NystromConfig
+from .ncut_nystrom import nystrom_propagate
+from .ncut_nystrom import _plain_ncut
+
+
+#TODO: automatically optimize click_weight based on the iou of fg and bg
+def ncut_click_prompt(
+        X: torch.Tensor,
+        fg_indices: np.ndarray,
+        bg_indices: np.ndarray = None,
+        click_weight: float = 0.5,
+        bg_weight: float = 0.1,
+        n_eig: int = 2,
+        quantile_sigma: float = 0.25,
+        device: str = None,
+        sigma: float = None,
+        affinity_fn: Callable[[torch.Tensor, torch.Tensor, float], torch.Tensor] = rbf_affinity,
+        exact_gradient: bool = False,
+        no_propagation: bool = False,
+        return_indices_and_sigma: bool = False,
+        **kwargs,
+) -> Union[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]]:
+
+    config = NystromConfig()
+    config.update(kwargs)
+
+    # use GPU if available
+    device = auto_device(X.device, device)
+
+    if bg_indices is None:
+        bg_indices = np.array([], dtype=np.int64)
+
+    # subsample for nystrom approximation
+    nystrom_indices = farthest_point_sampling(X, n_sample=config.n_sample, device=device)
+    nystrom_indices = torch.tensor(nystrom_indices, dtype=torch.long)
+    # remove fg and bg from fps_idx
+    nystrom_indices = nystrom_indices[~np.isin(nystrom_indices, np.concatenate([fg_indices, bg_indices]))]
+    # add fg and bg to fps_idx
+    nystrom_indices = np.concatenate([fg_indices, bg_indices, nystrom_indices])
+    fg_indices = np.arange(len(fg_indices))
+    bg_indices = np.arange(len(bg_indices)) + len(fg_indices)
+    n_fgbg = len(fg_indices) + len(bg_indices)
+    
+    nystrom_X = X[nystrom_indices].to(device)
+    
+    # find optimal sigma for affinity matrix
+    if sigma is None and affinity_fn == rbf_affinity:
+        sigma = find_sigma_by_degree(nystrom_X, quantile_sigma, affinity_fn)
+        # TODO: change to std()
+    elif sigma is None and affinity_fn == cosine_affinity:
+        sigma = 0.5
+
+    # compute Ncut on the nystrom sampled subgraph
+    A = affinity_fn(nystrom_X, sigma=sigma)
+    A = normalize_affinity(A)
+
+    # modify the affinity from the clicks
+    X_click = 1 * A[fg_indices].mean(0)
+    if len(bg_indices) > 0:
+        X_click = X_click - bg_weight * A[bg_indices].mean(0)
+    
+    X_click = X_click * A.shape[0]
+
+    A_click = affinity_fn(X_click.unsqueeze(1), sigma=0.5)
+    A_click = normalize_affinity(A_click)
+    
+    _A = click_weight * A_click + (1 - click_weight) * A
+        
+    nystrom_eigvec, eigval = _plain_ncut(_A, n_eig, exact_gradient=exact_gradient)
+    
+    if no_propagation:
+        return nystrom_eigvec, eigval, nystrom_indices, sigma
+
+    # propagate eigenvectors from subgraph to full graph
+    eigvec, nystrom_indices2 = nystrom_propagate(
+        nystrom_eigvec,
+        X,
+        nystrom_X,
+        n_neighbors=config.n_neighbors,
+        n_sample=config.n_sample2,
+        matmul_chunk_size=config.matmul_chunk_size,
+        device=device,
+        return_indices=True,
+    )
+    
+
+    if return_indices_and_sigma:
+        indices = nystrom_indices[nystrom_indices2]
+        return eigvec, eigval, indices, sigma
+
+    return eigvec, eigval

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/ncut_pytorch/ncuts/ncut_nystrom.py

@@ -6,10 +6,9 @@ import torch
 import numpy as np
 from ncut_pytorch.utils.sigma import find_sigma_by_degree
 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
+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
-from ncut_pytorch.utils.grad import grad_manager
 
 
 class NystromConfig:
@@ -23,7 +22,6 @@ class NystromConfig:
     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
-    move_output_to_cpu = True       # if True, will move output to cpu, saves VRAM
     
     def update(self, kwargs: dict):
         for key, value in kwargs.items():
@@ -36,15 +34,16 @@ class NystromConfig:
 def ncut_fn(
         X: torch.Tensor,
         n_eig: int = 100,
-        track_grad: bool = False,
-        d_sigma: float = None,
-        device: str = None,
-        sigma: float = None,
-        repulsion_sigma: float = None,
-        repulsion_weight: float = 0.2,
+        quantile_sigma: float = 0.25,
+        quantile_sigma_repulsion: float = 0.20,
+        sigma: float | None = None,
+        repulsion_sigma: float | None = None,
+        repulsion_weight: float | None = None,
+        affinity_fn: Union["rbf_affinity", "cosine_affinity"] = rbf_affinity,
         extrapolation_factor: float = 1.0,
+        exact_gradient: bool = False,
+        device: str | None = None,
         make_orthogonal: bool = False,
-        affinity_fn: Union["rbf_affinity", "cosine_affinity"] = rbf_affinity,
         no_propagation: bool = False,
         **kwargs,
 ) -> Union[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]]:
@@ -53,15 +52,15 @@ def ncut_fn(
     Args:
         X (torch.Tensor): input features, shape (N, D)
         n_eig (int): number of eigenvectors
-        track_grad (bool): keep track of pytorch gradients
-        d_sigma (float): affinity sigma parameter, lower d_sigma results in sharper eigenvectors
-        device (str): device, default 'auto' (auto detect GPU)
+        quantile_sigma (float): quantile of affinity sigma parameter, lower quantile_sigma results in sharper eigenvectors
+        quantile_sigma_repulsion (float): quantile of repulsion sigma parameter, lower quantile_sigma_repulsion results in sharper eigenvectors
         sigma (float): affinity parameter, override d_sigma if provided
         repulsion_sigma (float): (if use repulsion) repulsion sigma parameter, default None (no repulsion)
         repulsion_weight (float): (if use repulsion) repulsion weight, default 0.2
+        affinity_fn (callable): affinity function, default rbf_affinity. Should accept (X1, X2=None, sigma=float) and return affinity matrix
         extrapolation_factor (float): control how far can we extrapolate, larger extrapolation_factor means we can extrapolate further, default 1.0
+        exact_gradient (bool): use full spectrum and exact gradient, can be slower and unstable, default False
         make_orthogonal (bool): make eigenvectors orthogonal
-        affinity_fn (callable): affinity function, default rbf_affinity. Should accept (X1, X2=None, sigma=float) and return affinity matrix
         
     Returns:
         eigenvectors (torch.Tensor): shape (N, n_eig)
@@ -76,60 +75,67 @@ def ncut_fn(
     """
     config = NystromConfig()
     config.update(kwargs)
-
-    # use GPU if available
     device = auto_device(X.device, device)
 
-    # check if enough data for nystrom approximation
+    # 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])
+    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)
+
+    if repulsion_sigma and repulsion_weight:
+        nystrom_eigvec, eigval = ncut_with_repulsion(nystrom_X, n_eig, sigma, 
+            repulsion_sigma, repulsion_weight, affinity_fn, exact_gradient)
+    else:
+        A = affinity_fn(nystrom_X, sigma=sigma)
+        nystrom_eigvec, eigval = _plain_ncut(A, n_eig, exact_gradient)
+
+    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,
+        X,
+        nystrom_X,
+        extrapolation_factor=extrapolation_factor,
+        n_neighbors=config.n_neighbors,
+        n_sample=config.n_sample2,
+        matmul_chunk_size=config.matmul_chunk_size,
+        device=device,
+    )
+
+    # post-hoc orthogonalization
+    if make_orthogonal:
+        eigvec = gram_schmidt(eigvec)
 
-    with grad_manager(track_grad):
-        # subsample for nystrom approximation
-        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])
-        nystrom_X = X[nystrom_indices].to(device)
-
-        # find optimal sigma for affinity matrix
-        if sigma is None:
-            if affinity_fn == rbf_affinity:
-                sigma = find_sigma_by_degree(nystrom_X, d_sigma, affinity_fn)
-            elif affinity_fn == cosine_affinity:
-                sigma = 0.5
-            else:
-                raise ValueError(f"`sigma` needs to be provided for affinity function {affinity_fn}, (sigma=0.5)")
-                
-        if repulsion_sigma is not None:
-            nystrom_eigvec, eigval = ncut_with_repulsion(nystrom_X, n_eig, sigma_attraction=sigma, sigma_repulsion=repulsion_sigma, repulsion_weight=repulsion_weight, affinity_fn=affinity_fn)
-        else:
-            A = affinity_fn(nystrom_X, sigma=sigma)
-            nystrom_eigvec, eigval = _plain_ncut(A, n_eig)
-
-        if no_propagation:
-            return nystrom_eigvec, eigval, nystrom_indices, sigma
-
-        if not is_enough_data:
-            return nystrom_eigvec, eigval
-
-        # propagate eigenvectors from subgraph to full graph
-        eigvec = nystrom_propagate(
-            nystrom_eigvec,
-            X,
-            nystrom_X,
-            extrapolation_factor=extrapolation_factor,
-            n_neighbors=config.n_neighbors,
-            n_sample=config.n_sample2,
-            matmul_chunk_size=config.matmul_chunk_size,
-            device=device,
-            move_output_to_cpu=config.move_output_to_cpu,
-            track_grad=track_grad,
-        )
-
-        # post-hoc orthogonalization
-        if make_orthogonal:
-            eigvec = gram_schmidt(eigvec)
-
-        return eigvec, eigval
+    return eigvec, eigval
     
+def find_optimal_sigma(
+    X: torch.Tensor,
+    quantile_sigma: float = 0.25,
+    quantile_sigma_repulsion: float = 0.20,
+    sigma: float | None = None,
+    repulsion_sigma: float | None = None,
+    affinity_fn: Union["rbf_affinity", "cosine_affinity"] = rbf_affinity,
+):
+    """Find optimal sigma for affinity matrix and repulsion matrix."""
+    if affinity_fn == rbf_affinity:
+        sigma = sigma or find_sigma_by_degree(X, quantile_sigma, affinity_fn)
+        repulsion_sigma = repulsion_sigma or find_sigma_by_degree(X, quantile_sigma_repulsion, affinity_fn, init_sigma=sigma)
+    elif affinity_fn == cosine_affinity:
+        sigma = sigma or 0.5
+        repulsion_sigma = repulsion_sigma or 0.3
+    else:
+        if sigma is None:
+            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,
@@ -138,6 +144,7 @@ def ncut_with_repulsion(
     sigma_repulsion: float = None,
     repulsion_weight: float = 0.2,
     affinity_fn: Union["rbf_affinity", "cosine_affinity"] = cosine_affinity,
+    exact_gradient: bool = False,
     eps: float = 1e-8,
 ):
     A = affinity_fn(X, sigma=sigma_attraction)
@@ -148,7 +155,10 @@ def ncut_with_repulsion(
     D = D_A + D_R
     W = A - R + torch.diag(D_R)
     W = W / D[:, None]
-    eigvec, eigval, _ = grad_safe_eig_solve(W, n_eig)
+    if exact_gradient:
+        eigvec, eigval, _ = grad_safe_eig_solve(W, n_eig)
+    else:
+        eigvec, eigval, _ = svd_lowrank(W, n_eig)
     eigvec = correct_rotation(eigvec)
     return eigvec, eigval
 
@@ -156,9 +166,13 @@ def ncut_with_repulsion(
 def _plain_ncut(
         A: torch.Tensor,
         n_eig: int = 100,
+        exact_gradient: bool = False,
 ):
     A = normalize_affinity(A)
-    eigvec, eigval, _ = grad_safe_eig_solve(A, n_eig)
+    if exact_gradient:
+        eigvec, eigval, _ = grad_safe_eig_solve(A, n_eig)
+    else:
+        eigvec, eigval, _ = svd_lowrank(A, n_eig)
     eigvec = eigvec[:, :n_eig]
     eigval = eigval[:n_eig]
     eigvec = correct_rotation(eigvec)
@@ -170,7 +184,6 @@ def nystrom_propagate(
         X: torch.Tensor,
         nystrom_X: torch.Tensor,
         extrapolation_factor: float = 1.0,
-        track_grad: bool = False,        
         device: str = None,
         return_indices: bool = False,
         **kwargs,
@@ -183,7 +196,6 @@ def nystrom_propagate(
         X (torch.Tensor): input features for all nodes, shape (N, D)
         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
-        track_grad (bool): keep track of pytorch gradients, default False
         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
 
@@ -194,45 +206,44 @@ def nystrom_propagate(
     config = NystromConfig()
     config.update(kwargs)
 
-    with grad_manager(track_grad):
-        device = auto_device(nystrom_out.device, device)
-        indices = farthest_point_sampling(nystrom_out, config.n_sample2, device=device)
-        nystrom_out = nystrom_out[indices].to(device)
-        nystrom_X = nystrom_X[indices].to(device)
-        
-        sigma = find_sigma_by_degree(nystrom_X, affinity_fn=rbf_affinity)
-        sigma = sigma * extrapolation_factor
-        
-        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)
-        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)
-
-            if config.move_output_to_cpu and not track_grad:
-                out = out.to("cpu")
-            all_outs.append(out)
-
-        all_outs = torch.cat(all_outs, dim=0)
-
-        if return_indices:
-            return all_outs, indices
-        
-        return all_outs
+    device = auto_device(nystrom_out.device, device)
+    output_device = X.device
+    indices = farthest_point_sampling(nystrom_out, config.n_sample2, device=device)
+    nystrom_out = nystrom_out[indices].to(device)
+    nystrom_X = nystrom_X[indices].to(device)
+    
+    sigma = find_sigma_by_degree(nystrom_X, affinity_fn=rbf_affinity, quantile_sigma=0.25)
+    sigma = sigma * extrapolation_factor
+    
+    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)
+    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)
+
+    all_outs = torch.cat(all_outs, dim=0)
+
+    if return_indices:
+        return all_outs, indices
+    
+    return all_outs
 
 
 

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/ncut_pytorch/utils/grad.py

@@ -1,8 +1,6 @@
-__all__ = ["rbf_eigvec_manual_grad", "grad_manager"]
+__all__ = ["rbf_eigvec_manual_grad"]
 
 import torch
-from contextlib import contextmanager
-
 
 @torch.no_grad()
 def rbf_eigvec_manual_grad(
@@ -115,40 +113,3 @@ def rbf_eigvec_manual_grad(
 
     return grad_u
 
-
-
-@contextmanager
-def grad_manager(enabled: bool):
-    """Context manager to temporarily set gradient computation mode.
-    
-    This context manager allows you to control gradient computation for a block
-    of code, and automatically restores the previous gradient state when exiting
-    the context.
-    
-    Args:
-        enabled (bool): If True, enables gradient tracking within the context.
-                        If False, disables gradient tracking within the context.
-    
-    Yields:
-        None
-        
-    Examples:
-        >>> import torch
-        >>> from ncut_pytorch.utils.grad import set_grad_enabled
-        >>> 
-        >>> # Disable gradients for inference
-        >>> with set_grad_enabled(False):
-        ...     result = model(input_tensor)
-        >>> 
-        >>> # Enable gradients for training
-        >>> with set_grad_enabled(True):
-        ...     loss = criterion(model(input_tensor), target)
-        ...     loss.backward()
-    """
-    prev_grad_state = torch.is_grad_enabled()
-    torch.set_grad_enabled(enabled)
-    try:
-        yield
-    finally:
-        torch.set_grad_enabled(prev_grad_state)
-

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/ncut_pytorch/utils/math.py

@@ -18,7 +18,7 @@ import logging
 import numpy as np
 import torch
 
-from .torch_mod import svd_lowrank
+from .torch_mod import svd_lowrank as my_svd_lowrank
 
 
 def check_gamma_deprecated(gamma: float | None) -> float:
@@ -122,6 +122,23 @@ def pca_lowrank(
     return u @ torch.diag(s)
 
 
+def svd_lowrank(mat: torch.Tensor, q: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """SVD lowrank implementation for float16 and bfloat16."""
+    dtype = mat.dtype
+    try:
+        with torch.autocast(device_type=mat.device.type, enabled=False):
+            if dtype == torch.float16 or dtype == torch.bfloat16:
+                mat = mat.float()  # svd_lowrank does not support float16
+            u, s, v = my_svd_lowrank(mat, q=q + 10)
+    except RuntimeError:
+        if dtype == torch.float16 or dtype == torch.bfloat16:
+            mat = mat.float()
+        u, s, v = my_svd_lowrank(mat, q=q + 10)
+
+    u, s, v = u[:, :q], s[:q], v[:, :q]
+    return u.to(dtype), s.to(dtype), v.to(dtype)
+
+
 def quantile_min_max(
     x: torch.Tensor,
     q1: float = 0.01,

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/ncut_pytorch/utils/sigma.py

@@ -9,7 +9,7 @@ from .sample import farthest_point_sampling
 @torch.no_grad()
 def _find_sigma_by_degree(
     X: torch.Tensor,                    # [n_samples, n_features]
-    d_sigma: float | str | None = 'auto',
+    quantile_sigma: float = 0.25,
     affinity_fn: callable = rbf_affinity,
     X2: torch.Tensor | None = None,
     init_sigma: float = 0.5,
@@ -17,27 +17,21 @@ def _find_sigma_by_degree(
     max_iter: int = 100,
 ) -> float:
     """Binary search for optimal sigma to achieve target mean edge weight."""
-    if isinstance(d_sigma, float):
-        assert d_sigma > 0, "d_sigma must be positive"
+    if quantile_sigma <= 0 or quantile_sigma >= 1:
+        raise ValueError(f"quantile_sigma must be between 0 and 1, got {quantile_sigma}")
     sigma = init_sigma
     
-    # Find d_sigma if 'auto'
-    if d_sigma in ('auto', None):
-        scale_inv_sigma = sigma * X.std(0).sum()
-        current_degrees = affinity_fn(X, X2=X2, sigma=scale_inv_sigma).mean(1)
-        for _ in range(2):
-            current_degree = current_degrees.mean().item()
-            mask = current_degrees < current_degree
-            current_degrees = current_degrees[mask]
-        d_sigma = current_degrees.mean().item()
+    scale_inv_sigma = X.std(0).sum()
+    current_degrees = affinity_fn(X, X2=X2, sigma=scale_inv_sigma).mean(1)
+    target_degree = current_degrees.float().quantile(quantile_sigma).item()
     
     # Binary search for sigma
     current_degree = affinity_fn(X, X2=X2, sigma=sigma).mean().item()
     low, high = 0, float('inf')
-    tol = r_tol * d_sigma
+    tol = r_tol * target_degree
     i_iter = 0
-    while abs(current_degree - d_sigma) > tol and i_iter < max_iter:
-        if current_degree > d_sigma:
+    while abs(current_degree - target_degree) > tol and i_iter < max_iter:
+        if current_degree > target_degree:
             high = sigma
             sigma = (low + sigma) / 2
         else:
@@ -52,7 +46,7 @@ def _find_sigma_by_degree(
 @torch.no_grad()
 def find_sigma_by_degree(
     X: torch.Tensor,                    # [n_samples, n_features]
-    d_sigma: float | str | None = 'auto',
+    quantile_sigma: float = 0.25,
     affinity_fn: callable = rbf_affinity,
     X2: torch.Tensor | None = None,
     init_sigma: float = 0.5,
@@ -62,4 +56,4 @@ def find_sigma_by_degree(
 ) -> float:
     """Find sigma after FPS-based downsampling for efficiency."""
     indices = farthest_point_sampling(X, n_sample)
-    return _find_sigma_by_degree(X[indices], d_sigma, affinity_fn, X2=X2, init_sigma=init_sigma, r_tol=r_tol, max_iter=max_iter)
+    return _find_sigma_by_degree(X[indices], quantile_sigma, affinity_fn, X2=X2, init_sigma=init_sigma, r_tol=r_tol, max_iter=max_iter)

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/ncut_pytorch.egg-info/PKG-INFO

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

{ncut_pytorch-3.0.0.dev0 → ncut_pytorch-3.0.0.dev2}/pyproject.toml

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

ncut_pytorch-3.0.0.dev0/ncut_pytorch/ncuts/ncut_click.py

@@ -1,106 +0,0 @@
-__all__ = ['ncut_click_prompt']
-
-from typing import Callable, Union
-
-import numpy as np
-import torch
-
-from ncut_pytorch.utils.sigma import find_sigma_by_degree
-from ncut_pytorch.utils.math import rbf_affinity, cosine_affinity, normalize_affinity
-from ncut_pytorch.utils.sample import farthest_point_sampling
-from ncut_pytorch.utils.device import auto_device
-from ncut_pytorch.utils.grad import grad_manager
-from .ncut_nystrom import NystromConfig
-from .ncut_nystrom import nystrom_propagate
-from .ncut_nystrom import _plain_ncut
-
-
-#TODO: automatically optimize click_weight based on the iou of fg and bg
-def ncut_click_prompt(
-        X: torch.Tensor,
-        fg_indices: np.ndarray,
-        bg_indices: np.ndarray = None,
-        click_weight: float = 0.5,
-        bg_weight: float = 0.1,
-        n_eig: int = 2,
-        track_grad: bool = False,
-        d_sigma: float = None,
-        device: str = None,
-        sigma: float = None,
-        affinity_fn: Callable[[torch.Tensor, torch.Tensor, float], torch.Tensor] = rbf_affinity,
-        no_propagation: bool = False,
-        return_indices_and_sigma: bool = False,
-        **kwargs,
-) -> Union[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]]:
-
-    config = NystromConfig()
-    config.update(kwargs)
-
-    # use GPU if available
-    device = auto_device(X.device, device)
-
-    with grad_manager(track_grad):
-        if bg_indices is None:
-            bg_indices = np.array([], dtype=np.int64)
-
-        # subsample for nystrom approximation
-        nystrom_indices = farthest_point_sampling(X, n_sample=config.n_sample, device=device)
-        nystrom_indices = torch.tensor(nystrom_indices, dtype=torch.long)
-        # remove fg and bg from fps_idx
-        nystrom_indices = nystrom_indices[~np.isin(nystrom_indices, np.concatenate([fg_indices, bg_indices]))]
-        # add fg and bg to fps_idx
-        nystrom_indices = np.concatenate([fg_indices, bg_indices, nystrom_indices])
-        fg_indices = np.arange(len(fg_indices))
-        bg_indices = np.arange(len(bg_indices)) + len(fg_indices)
-        n_fgbg = len(fg_indices) + len(bg_indices)
-        
-        nystrom_X = X[nystrom_indices].to(device)
-        
-        # find optimal sigma for affinity matrix
-        if sigma is None and affinity_fn == rbf_affinity:
-            sigma = find_sigma_by_degree(nystrom_X, d_sigma, affinity_fn)
-            # TODO: change to std()
-        elif sigma is None and affinity_fn == cosine_affinity:
-            sigma = 0.5
-
-        # compute Ncut on the nystrom sampled subgraph
-        A = affinity_fn(nystrom_X, sigma=sigma)
-        A = normalize_affinity(A)
-
-        # modify the affinity from the clicks
-        X_click = 1 * A[fg_indices].mean(0)
-        if len(bg_indices) > 0:
-            X_click = X_click - bg_weight * A[bg_indices].mean(0)
-        
-        X_click = X_click * A.shape[0]
-
-        A_click = affinity_fn(X_click.unsqueeze(1), sigma=0.5)
-        A_click = normalize_affinity(A_click)
-        
-        _A = click_weight * A_click + (1 - click_weight) * A
-            
-        nystrom_eigvec, eigval = _plain_ncut(_A, n_eig)
-        
-        if no_propagation:
-            return nystrom_eigvec, eigval, nystrom_indices, sigma
-
-        # propagate eigenvectors from subgraph to full graph
-        eigvec, nystrom_indices2 = nystrom_propagate(
-            nystrom_eigvec,
-            X,
-            nystrom_X,
-            n_neighbors=config.n_neighbors,
-            n_sample=config.n_sample2,
-            matmul_chunk_size=config.matmul_chunk_size,
-            device=device,
-            move_output_to_cpu=config.move_output_to_cpu,
-            track_grad=track_grad,
-            return_indices=True,
-        )
-        
-
-        if return_indices_and_sigma:
-            indices = nystrom_indices[nystrom_indices2]
-            return eigvec, eigval, indices, sigma
-
-        return eigvec, eigval