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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ncut_pytorch
-Version: 3.0.0.dev0
+Version: 3.0.0.dev1
 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.dev1}/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.dev1}/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.dev1/ncut_pytorch/ncuts/ncut_click.py

@@ -0,0 +1,101 @@
+__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,
+        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)
+
+    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,
+        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.dev1}/ncut_pytorch/ncuts/ncut_nystrom.py

@@ -9,7 +9,6 @@ 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.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,7 +34,6 @@ 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,
@@ -53,7 +50,6 @@ 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)
         sigma (float): affinity parameter, override d_sigma if provided
@@ -83,52 +79,49 @@ def ncut_fn(
     # check if enough data for nystrom approximation
     is_enough_data = X.shape[0] > config.n_sample
 
-    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)
+    # 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:
-            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
+            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,
+    )
+
+    # post-hoc orthogonalization
+    if make_orthogonal:
+        eigvec = gram_schmidt(eigvec)
+
+    return eigvec, eigval
     
 
 def ncut_with_repulsion(
@@ -170,7 +163,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 +175,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 +185,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)
+    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.dev1}/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.dev1}/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.dev1
 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.dev1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ncut_pytorch"
-version = "3.0.0dev0"
+version = "3.0.0dev1"
 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