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

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

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

@@ -22,10 +22,11 @@ def ncut_click_prompt(
         click_weight: float = 0.5,
         bg_weight: float = 0.1,
         n_eig: int = 2,
-        d_sigma: float = None,
+        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,
@@ -55,7 +56,7 @@ def ncut_click_prompt(
     
     # 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)
+        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
@@ -76,7 +77,7 @@ def ncut_click_prompt(
     
     _A = click_weight * A_click + (1 - click_weight) * A
         
-    nystrom_eigvec, eigval = _plain_ncut(_A, n_eig)
+    nystrom_eigvec, eigval = _plain_ncut(_A, n_eig, exact_gradient=exact_gradient)
     
     if no_propagation:
         return nystrom_eigvec, eigval, nystrom_indices, sigma

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

@@ -6,7 +6,7 @@ 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
 
@@ -34,14 +34,16 @@ class NystromConfig:
 def ncut_fn(
         X: torch.Tensor,
         n_eig: int = 100,
-        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]]:
@@ -50,14 +52,15 @@ def ncut_fn(
     Args:
         X (torch.Tensor): input features, shape (N, D)
         n_eig (int): number of eigenvectors
-        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)
@@ -72,37 +75,28 @@ 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
-    is_enough_data = X.shape[0] > config.n_sample
-
     # 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)
 
-    # 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)
+    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)
+        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
@@ -123,6 +117,25 @@ def ncut_fn(
 
     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,
@@ -131,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)
@@ -141,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
 
@@ -149,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)
@@ -191,7 +212,7 @@ def nystrom_propagate(
     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 = 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)

{ncut_pytorch-3.0.0.dev1 → 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.dev1 → 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.dev1 → 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.dev1
+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.dev1 → ncut_pytorch-3.0.0.dev2}/pyproject.toml

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