nystrom-ncut 0.0.2__tar.gz → 0.0.3__tar.gz

{nystrom_ncut-0.0.2/src/nystrom_ncut.egg-info → nystrom_ncut-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nystrom_ncut
-Version: 0.0.2
+Version: 0.0.3
 Summary: Normalized Cut and Nyström Approximation
 Author-email: Huzheng Yang <huze.yann@gmail.com>, Wentinn Liao <wentinn.liao@gmail.com>
 Project-URL: Documentation, https://github.com/JophiArcana/Nystrom-NCUT/

{nystrom_ncut-0.0.2 → nystrom_ncut-0.0.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nystrom_ncut"
-version = "0.0.2"
+version = "0.0.3"
 authors = [
     { name = "Huzheng Yang", email = "huze.yann@gmail.com" },
     { name = "Wentinn Liao", email = "wentinn.liao@gmail.com" },

{nystrom_ncut-0.0.2 → nystrom_ncut-0.0.3}/src/nystrom_ncut/ncut_pytorch.py

@@ -94,20 +94,19 @@ class NCUT(OnlineNystrom):
 
     def __init__(
         self,
-        num_eig: int = 100,
+        n_components: int = 100,
         affinity_focal_gamma: float = 1.0,
         num_sample: int = 10000,
         sample_method: Literal["farthest", "random"] = "farthest",
         distance: DistanceOptions = "cosine",
         eig_solver: EigSolverOptions = "svd_lowrank",
         normalize_features: bool = None,
-        device: str = None,
         move_output_to_cpu: bool = False,
-        matmul_chunk_size: int = 8096,
+        chunk_size: int = 8192,
     ):
         """
         Args:
-            num_eig (int): number of top eigenvectors to return
+            n_components (int): number of top eigenvectors to return
             affinity_focal_gamma (float): affinity matrix temperature, lower t reduce the not-so-connected edge weights,
                 smaller t result in more sharp eigenvectors.
             num_sample (int): number of samples for Nystrom-like approximation,
@@ -118,17 +117,15 @@ class NCUT(OnlineNystrom):
             eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh'].
             normalize_features (bool): normalize input features before computing affinity matrix,
                 default 'None' is True for cosine distance, False for euclidean distance and rbf
-            device (str): device to use for eigen computation,
-                move to GPU to speeds up a bit (~5x faster)
             move_output_to_cpu (bool): move output to CPU, set to True if you have memory issue
-            matmul_chunk_size (int): chunk size for large-scale matrix multiplication
+            chunk_size (int): chunk size for large-scale matrix multiplication
         """
         OnlineNystrom.__init__(
             self,
-            n_components=num_eig,
+            n_components=n_components,
             kernel=LaplacianKernel(affinity_focal_gamma, distance, eig_solver),
             eig_solver=eig_solver,
-            chunk_size=matmul_chunk_size,
+            chunk_size=chunk_size,
         )
         self.num_sample = num_sample
         self.sample_method = sample_method
@@ -140,19 +137,14 @@ class NCUT(OnlineNystrom):
             if distance in ["euclidean", "rbf"]:
                 self.normalize_features = False
 
-        self.device = device
         self.move_output_to_cpu = move_output_to_cpu
-        self.matmul_chunk_size = matmul_chunk_size
+        self.chunk_size = chunk_size
 
     def _fit_helper(
         self,
         features: torch.Tensor,
         precomputed_sampled_indices: torch.Tensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # move subgraph gpu to speed up
-        original_device = features.device
-        device = original_device if self.device is None else self.device
-
         _n = features.shape[0]
         if self.num_sample >= _n:
             logging.info(
@@ -184,13 +176,13 @@ class NCUT(OnlineNystrom):
                 num_sample=self.num_sample,
                 sample_method=self.sample_method,
             )
-        sampled_features = features[sampled_indices].to(device)
+        sampled_features = features[sampled_indices]
         OnlineNystrom.fit(self, sampled_features)
 
         _n_not_sampled = _n - len(sampled_features)
         if _n_not_sampled > 0:
-            unsampled_indices = torch.full((_n,), True).scatter(0, sampled_indices, False)
-            unsampled_features = features[unsampled_indices].to(device)
+            unsampled_indices = torch.full((_n,), True, device=features.device).scatter_(0, sampled_indices, False)
+            unsampled_features = features[unsampled_indices]
             V_unsampled, _ = OnlineNystrom.update(self, unsampled_features)
         else:
             unsampled_indices = V_unsampled = None
@@ -231,7 +223,7 @@ class NCUT(OnlineNystrom):
         V_sampled, L = OnlineNystrom.transform(self)
 
         if unsampled_indices is not None:
-            V = torch.zeros((len(unsampled_indices), self.n_components))
+            V = torch.zeros((len(unsampled_indices), self.n_components), device=features.device)
             V[~unsampled_indices] = V_sampled
             V[unsampled_indices] = V_unsampled
         else:

{nystrom_ncut-0.0.2 → nystrom_ncut-0.0.3}/src/nystrom_ncut/propagation_utils.py

@@ -43,7 +43,7 @@ def run_subgraph_sampling(
             sampled_indices = torch.randperm(features.shape[0])[:num_sample]
         else:
             raise ValueError("sample_method should be 'farthest' or 'random'")
-    return sampled_indices
+    return sampled_indices.to(features.device)
 
 
 def farthest_point_sampling(
@@ -139,7 +139,7 @@ def propagate_knn(
     knn: int = 10,
     distance: Literal["cosine", "euclidean", "rbf"] = "cosine",
     affinity_focal_gamma: float = 1.0,
-    chunk_size: int = 8096,
+    chunk_size: int = 8192,
     device: str = None,
     move_output_to_cpu: bool = False,
 ):
@@ -206,7 +206,7 @@ def propagate_nearest(
     inp_features: torch.Tensor,
     subgraph_features: torch.Tensor,
     distance: Literal["cosine", "euclidean", "rbf"] = "cosine",
-    chunk_size: int = 8096,
+    chunk_size: int = 8192,
     device: str = None,
     move_output_to_cpu: bool = False,
 ):
@@ -254,7 +254,7 @@ def propagate_eigenvectors(
         knn (int): number of KNN to propagate eigenvectors, default 3
         num_sample (int): number of samples for subgraph sampling, default 50000
         sample_method (str): sample method, 'farthest' (default) or 'random'
-        chunk_size (int): chunk size for matrix multiplication, default 8096
+        chunk_size (int): chunk size for matrix multiplication, default 8192
         device (str): device to use for computation, if None, will not change device
     Returns:
         torch.Tensor: propagated eigenvectors, shape (n_new_samples, num_eig)

{nystrom_ncut-0.0.2 → nystrom_ncut-0.0.3}/src/nystrom_ncut/visualize_utils.py

@@ -420,7 +420,7 @@ def propagate_rgb_color(
     knn: int = 10,
     num_sample: int = 1000,
     sample_method: Literal["farthest", "random"] = "farthest",
-    chunk_size: int = 8096,
+    chunk_size: int = 8192,
     device: str = None,
 ):
     """Propagate RGB color to new nodes using KNN.
@@ -431,7 +431,7 @@ def propagate_rgb_color(
         knn (int): number of KNN to propagate RGB color, default 1
         num_sample (int): number of samples for subgraph sampling, default 50000
         sample_method (str): sample method, 'farthest' (default) or 'random'
-        chunk_size (int): chunk size for matrix multiplication, default 8096
+        chunk_size (int): chunk size for matrix multiplication, default 8192
         device (str): device to use for computation, if None, will not change device
     Returns:
         torch.Tensor: propagated RGB color for each data sample, shape (n_new_samples, 3)

{nystrom_ncut-0.0.2 → nystrom_ncut-0.0.3/src/nystrom_ncut.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nystrom_ncut
-Version: 0.0.2
+Version: 0.0.3
 Summary: Normalized Cut and Nyström Approximation
 Author-email: Huzheng Yang <huze.yann@gmail.com>, Wentinn Liao <wentinn.liao@gmail.com>
 Project-URL: Documentation, https://github.com/JophiArcana/Nystrom-NCUT/