sped up Eig Decomposition w/ cupy 40 min => 40 secs

graphgps/loader/master_loader.py

@@ -203,6 +203,8 @@ def convert_to_int(ds, prop):
         # Estimate directedness based on 10 graphs to save time.
         is_undirected = all(d.is_undirected() for d in dataset[:10])
         logging.info(f"  ...estimated to be undirected: {is_undirected}")
+        msg = "GPU installed, using CUPY to preprocess, brrr" if torch.cuda.is_available() and cfg.prep_w_GPU else "using Numpy instead of Cupy w/ GPU"
+        logging.info(msg)
         pre_transform_in_memory(dataset,
                                 partial(compute_posenc_stats,
                                         pe_types=pe_enabled_list,
@@ -517,6 +519,7 @@ def preformat_Peptides(dataset_dir, name):
     if dataset_type == 'functional':
         dataset = PeptidesFunctionalDataset(dataset_dir)
     elif dataset_type == 'structural':
+#         dataset = PeptidesStructuralDataset(dataset_dir, pre_transform=partial(task_specific_preprocessing, cfg=cfg))
         dataset = PeptidesStructuralDataset(dataset_dir)
     s_dict = dataset.get_idx_split()
     dataset.split_idxs = [s_dict[s] for s in ['train', 'val', 'test']]

graphgps/transform/posenc_stats.py

@@ -52,13 +52,25 @@ def compute_posenc_stats(data, pe_types, is_undirected, cfg):
 
     # Eigen values and vectors.
     evals, evects = None, None
+    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    undir_edge_index = undir_edge_index.to(dev)
     if 'LapPE' in pe_types or 'EquivStableLapPE' in pe_types:
-        # Eigen-decomposition with numpy, can be reused for Heat kernels.
-        L = to_scipy_sparse_matrix(
-            *get_laplacian(undir_edge_index, normalization=laplacian_norm_type,
-                           num_nodes=N)
-        )
-        evals, evects = np.linalg.eigh(L.toarray())
+        try:
+            if not cfg.prep_w_GPU:
+                raise ImportError
+            import cupy as cp
+            edge_i, edge_w = get_laplacian(undir_edge_index, normalization=laplacian_norm_type, num_nodes=N)
+            dense_L = to_dense_adj(edge_index=edge_i, edge_attr=edge_w).squeeze() 
+            L = cp.asarray(dense_L)
+            evals, evects = cp.linalg.eigh(L)
+            evals, evects = cp.asnumpy(evals), cp.asnumpy(evects)
+        except ImportError:
+            # Eigen-decomposition with numpy, can be reused for Heat kernels.
+            L = to_scipy_sparse_matrix(
+                *get_laplacian(undir_edge_index, normalization=laplacian_norm_type,
+                               num_nodes=N)
+            )
+            evals, evects = np.linalg.eigh(L.toarray())
 
         if 'LapPE' in pe_types:
             max_freqs=cfg.posenc_LapPE.eigen.max_freqs
@@ -77,11 +89,21 @@ def compute_posenc_stats(data, pe_types, is_undirected, cfg):
         norm_type = cfg.posenc_SignNet.eigen.laplacian_norm.lower()
         if norm_type == 'none':
             norm_type = None
-        L = to_scipy_sparse_matrix(
-            *get_laplacian(undir_edge_index, normalization=norm_type,
-                           num_nodes=N)
-        )
-        evals_sn, evects_sn = np.linalg.eigh(L.toarray())
+        try:
+            if not cfg.prep_w_GPU:
+                raise ImportError
+            import cupy as cp
+            edge_i, edge_w = get_laplacian(undir_edge_index, normalization=laplacian_norm_type, num_nodes=N)
+            dense_L = to_dense_adj(edge_index=edge_i, edge_attr=edge_w).squeeze()
+            L = cp.asarray(dense_L)
+            evals, evects = cp.linalg.eigh(L)
+            evals_sn, evects_sn = cp.asnumpy(evals), cp.asnumpy(evects) 
+        except ImportError:
+            L = to_scipy_sparse_matrix(
+                *get_laplacian(undir_edge_index, normalization=norm_type, num_nodes=N)
+            )
+            evals_sn, evects_sn = np.linalg.eigh(L.toarray())
+
         data.eigvals_sn, data.eigvecs_sn = get_lap_decomp_stats(
             evals=evals_sn, evects=evects_sn,
             max_freqs=cfg.posenc_SignNet.eigen.max_freqs,
@@ -102,10 +124,21 @@ def compute_posenc_stats(data, pe_types, is_undirected, cfg):
         # Get the eigenvalues and eigenvectors of the regular Laplacian,
         # if they have not yet been computed for 'eigen'.
         if laplacian_norm_type is not None or evals is None or evects is None:
-            L_heat = to_scipy_sparse_matrix(
-                *get_laplacian(undir_edge_index, normalization=None, num_nodes=N)
-            )
-            evals_heat, evects_heat = np.linalg.eigh(L_heat.toarray())
+            ## normalization None for heat Kernels?  
+            try:
+                if not cfg.prep_w_GPU:
+                    raise ImportError
+                import cupy as cp
+                edge_i, edge_w = get_laplacian(undir_edge_index, normalization=None, num_nodes=N)
+                dense_L = to_dense_adj(edge_index=edge_i, edge_attr=edge_w).squeeze()
+                L_heat = cp.asarray(dense_L)
+                evals, evects = cp.linalg.eigh(L_heat)
+                evals_heat, evects_heat = cp.asnumpy(evals), cp.asnumpy(evects)
+            except ImportError:
+                L_heat = to_scipy_sparse_matrix(
+                    *get_laplacian(undir_edge_index, normalization=None, num_nodes=N)
+                )
+                evals_heat, evects_heat = np.linalg.eigh(L_heat.toarray())
         else:
             evals_heat, evects_heat = evals, evects
         evals_heat = torch.from_numpy(evals_heat)

main.py

@@ -116,6 +116,8 @@ def run_loop_settings():
 if __name__ == '__main__':
     # Load cmd line args
     args = parse_args()
+    # default prep w/ GPU
+    cfg.prep_w_GPU = True
     # Load config file
     set_cfg(cfg)
     load_cfg(cfg, args)