diff --git a/graphgps/loader/master_loader.py b/graphgps/loader/master_loader.py index edadb7e8..c9be2d46 100644 --- a/graphgps/loader/master_loader.py +++ b/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']] diff --git a/graphgps/transform/posenc_stats.py b/graphgps/transform/posenc_stats.py index 90ebcf74..afe9ca53 100644 --- a/graphgps/transform/posenc_stats.py +++ b/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) diff --git a/main.py b/main.py index 5044a591..5ded12a9 100644 --- a/main.py +++ b/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)