From d920d6a1db28143f8ea40460b7d65419ea633b55 Mon Sep 17 00:00:00 2001 From: euku Date: Sat, 2 Dec 2023 11:41:26 -0800 Subject: [PATCH 1/5] sped up Eig Decomposition w/ cupy 40 min => 40 secs --- graphgps/transform/posenc_stats.py | 41 ++++++++++++++++++------------ 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/graphgps/transform/posenc_stats.py b/graphgps/transform/posenc_stats.py index 90ebcf74..ff21452f 100644 --- a/graphgps/transform/posenc_stats.py +++ b/graphgps/transform/posenc_stats.py @@ -8,7 +8,7 @@ to_undirected, to_dense_adj, scatter) from torch_geometric.utils.num_nodes import maybe_num_nodes from graphgps.encoder.graphormer_encoder import graphormer_pre_processing - +import cupy as cp def compute_posenc_stats(data, pe_types, is_undirected, cfg): """Precompute positional encodings for the given graph. @@ -52,13 +52,14 @@ def compute_posenc_stats(data, pe_types, is_undirected, cfg): # Eigen values and vectors. evals, evects = None, None + dev = torch.device("cuda") 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()) + undir_edge_index = undir_edge_index.to(dev) + 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) if 'LapPE' in pe_types: max_freqs=cfg.posenc_LapPE.eigen.max_freqs @@ -77,11 +78,12 @@ 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()) + 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) + 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 +104,17 @@ 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()) + ## is normalization None here on purpose? + 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) + +# 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) From b136427c90b0e744c2f84610e18fb6be3318ffd1 Mon Sep 17 00:00:00 2001 From: Eugene29 <52263376+Eugene29@users.noreply.github.com> Date: Sun, 17 Dec 2023 16:51:57 -0800 Subject: [PATCH 2/5] Delete main.py --- main.py | 176 -------------------------------------------------------- 1 file changed, 176 deletions(-) delete mode 100644 main.py diff --git a/main.py b/main.py deleted file mode 100644 index 5044a591..00000000 --- a/main.py +++ /dev/null @@ -1,176 +0,0 @@ -import datetime -import os -import torch -import logging - -import graphgps # noqa, register custom modules -from graphgps.agg_runs import agg_runs -from graphgps.optimizer.extra_optimizers import ExtendedSchedulerConfig - -from torch_geometric.graphgym.cmd_args import parse_args -from torch_geometric.graphgym.config import (cfg, dump_cfg, - set_cfg, load_cfg, - makedirs_rm_exist) -from torch_geometric.graphgym.loader import create_loader -from torch_geometric.graphgym.logger import set_printing -from torch_geometric.graphgym.optim import create_optimizer, \ - create_scheduler, OptimizerConfig -from torch_geometric.graphgym.model_builder import create_model -from torch_geometric.graphgym.train import GraphGymDataModule, train -from torch_geometric.graphgym.utils.comp_budget import params_count -from torch_geometric.graphgym.utils.device import auto_select_device -from torch_geometric.graphgym.register import train_dict -from torch_geometric import seed_everything - -from graphgps.finetuning import load_pretrained_model_cfg, \ - init_model_from_pretrained -from graphgps.logger import create_logger - - -torch.backends.cuda.matmul.allow_tf32 = True # Default False in PyTorch 1.12+ -torch.backends.cudnn.allow_tf32 = True # Default True - - -def new_optimizer_config(cfg): - return OptimizerConfig(optimizer=cfg.optim.optimizer, - base_lr=cfg.optim.base_lr, - weight_decay=cfg.optim.weight_decay, - momentum=cfg.optim.momentum) - - -def new_scheduler_config(cfg): - return ExtendedSchedulerConfig( - scheduler=cfg.optim.scheduler, - steps=cfg.optim.steps, lr_decay=cfg.optim.lr_decay, - max_epoch=cfg.optim.max_epoch, reduce_factor=cfg.optim.reduce_factor, - schedule_patience=cfg.optim.schedule_patience, min_lr=cfg.optim.min_lr, - num_warmup_epochs=cfg.optim.num_warmup_epochs, - train_mode=cfg.train.mode, eval_period=cfg.train.eval_period) - - -def custom_set_out_dir(cfg, cfg_fname, name_tag): - """Set custom main output directory path to cfg. - Include the config filename and name_tag in the new :obj:`cfg.out_dir`. - - Args: - cfg (CfgNode): Configuration node - cfg_fname (string): Filename for the yaml format configuration file - name_tag (string): Additional name tag to identify this execution of the - configuration file, specified in :obj:`cfg.name_tag` - """ - run_name = os.path.splitext(os.path.basename(cfg_fname))[0] - run_name += f"-{name_tag}" if name_tag else "" - cfg.out_dir = os.path.join(cfg.out_dir, run_name) - - -def custom_set_run_dir(cfg, run_id): - """Custom output directory naming for each experiment run. - - Args: - cfg (CfgNode): Configuration node - run_id (int): Main for-loop iter id (the random seed or dataset split) - """ - cfg.run_dir = os.path.join(cfg.out_dir, str(run_id)) - # Make output directory - if cfg.train.auto_resume: - os.makedirs(cfg.run_dir, exist_ok=True) - else: - makedirs_rm_exist(cfg.run_dir) - - -def run_loop_settings(): - """Create main loop execution settings based on the current cfg. - - Configures the main execution loop to run in one of two modes: - 1. 'multi-seed' - Reproduces default behaviour of GraphGym when - args.repeats controls how many times the experiment run is repeated. - Each iteration is executed with a random seed set to an increment from - the previous one, starting at initial cfg.seed. - 2. 'multi-split' - Executes the experiment run over multiple dataset splits, - these can be multiple CV splits or multiple standard splits. The random - seed is reset to the initial cfg.seed value for each run iteration. - - Returns: - List of run IDs for each loop iteration - List of rng seeds to loop over - List of dataset split indices to loop over - """ - if len(cfg.run_multiple_splits) == 0: - # 'multi-seed' run mode - num_iterations = args.repeat - seeds = [cfg.seed + x for x in range(num_iterations)] - split_indices = [cfg.dataset.split_index] * num_iterations - run_ids = seeds - else: - # 'multi-split' run mode - if args.repeat != 1: - raise NotImplementedError("Running multiple repeats of multiple " - "splits in one run is not supported.") - num_iterations = len(cfg.run_multiple_splits) - seeds = [cfg.seed] * num_iterations - split_indices = cfg.run_multiple_splits - run_ids = split_indices - return run_ids, seeds, split_indices - - -if __name__ == '__main__': - # Load cmd line args - args = parse_args() - # Load config file - set_cfg(cfg) - load_cfg(cfg, args) - custom_set_out_dir(cfg, args.cfg_file, cfg.name_tag) - dump_cfg(cfg) - # Set Pytorch environment - torch.set_num_threads(cfg.num_threads) - # Repeat for multiple experiment runs - for run_id, seed, split_index in zip(*run_loop_settings()): - # Set configurations for each run - custom_set_run_dir(cfg, run_id) - set_printing() - cfg.dataset.split_index = split_index - cfg.seed = seed - cfg.run_id = run_id - seed_everything(cfg.seed) - auto_select_device() - if cfg.pretrained.dir: - cfg = load_pretrained_model_cfg(cfg) - logging.info(f"[*] Run ID {run_id}: seed={cfg.seed}, " - f"split_index={cfg.dataset.split_index}") - logging.info(f" Starting now: {datetime.datetime.now()}") - # Set machine learning pipeline - loaders = create_loader() - loggers = create_logger() - model = create_model() - if cfg.pretrained.dir: - model = init_model_from_pretrained( - model, cfg.pretrained.dir, cfg.pretrained.freeze_main, - cfg.pretrained.reset_prediction_head, seed=cfg.seed - ) - optimizer = create_optimizer(model.parameters(), - new_optimizer_config(cfg)) - scheduler = create_scheduler(optimizer, new_scheduler_config(cfg)) - # Print model info - logging.info(model) - logging.info(cfg) - cfg.params = params_count(model) - logging.info('Num parameters: %s', cfg.params) - # Start training - if cfg.train.mode == 'standard': - if cfg.wandb.use: - logging.warning("[W] WandB logging is not supported with the " - "default train.mode, set it to `custom`") - datamodule = GraphGymDataModule() - train(model, datamodule, logger=True) - else: - train_dict[cfg.train.mode](loggers, loaders, model, optimizer, - scheduler) - # Aggregate results from different seeds - try: - agg_runs(cfg.out_dir, cfg.metric_best) - except Exception as e: - logging.info(f"Failed when trying to aggregate multiple runs: {e}") - # When being launched in batch mode, mark a yaml as done - if args.mark_done: - os.rename(args.cfg_file, f'{args.cfg_file}_done') - logging.info(f"[*] All done: {datetime.datetime.now()}") From 2ef714cb4c710663c2436d8d986e6072531c56d7 Mon Sep 17 00:00:00 2001 From: Eugene29 <52263376+Eugene29@users.noreply.github.com> Date: Sun, 17 Dec 2023 16:52:19 -0800 Subject: [PATCH 3/5] Add files via upload device agnostic main --- main.py | 178 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 178 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 00000000..5ded12a9 --- /dev/null +++ b/main.py @@ -0,0 +1,178 @@ +import datetime +import os +import torch +import logging + +import graphgps # noqa, register custom modules +from graphgps.agg_runs import agg_runs +from graphgps.optimizer.extra_optimizers import ExtendedSchedulerConfig + +from torch_geometric.graphgym.cmd_args import parse_args +from torch_geometric.graphgym.config import (cfg, dump_cfg, + set_cfg, load_cfg, + makedirs_rm_exist) +from torch_geometric.graphgym.loader import create_loader +from torch_geometric.graphgym.logger import set_printing +from torch_geometric.graphgym.optim import create_optimizer, \ + create_scheduler, OptimizerConfig +from torch_geometric.graphgym.model_builder import create_model +from torch_geometric.graphgym.train import GraphGymDataModule, train +from torch_geometric.graphgym.utils.comp_budget import params_count +from torch_geometric.graphgym.utils.device import auto_select_device +from torch_geometric.graphgym.register import train_dict +from torch_geometric import seed_everything + +from graphgps.finetuning import load_pretrained_model_cfg, \ + init_model_from_pretrained +from graphgps.logger import create_logger + + +torch.backends.cuda.matmul.allow_tf32 = True # Default False in PyTorch 1.12+ +torch.backends.cudnn.allow_tf32 = True # Default True + + +def new_optimizer_config(cfg): + return OptimizerConfig(optimizer=cfg.optim.optimizer, + base_lr=cfg.optim.base_lr, + weight_decay=cfg.optim.weight_decay, + momentum=cfg.optim.momentum) + + +def new_scheduler_config(cfg): + return ExtendedSchedulerConfig( + scheduler=cfg.optim.scheduler, + steps=cfg.optim.steps, lr_decay=cfg.optim.lr_decay, + max_epoch=cfg.optim.max_epoch, reduce_factor=cfg.optim.reduce_factor, + schedule_patience=cfg.optim.schedule_patience, min_lr=cfg.optim.min_lr, + num_warmup_epochs=cfg.optim.num_warmup_epochs, + train_mode=cfg.train.mode, eval_period=cfg.train.eval_period) + + +def custom_set_out_dir(cfg, cfg_fname, name_tag): + """Set custom main output directory path to cfg. + Include the config filename and name_tag in the new :obj:`cfg.out_dir`. + + Args: + cfg (CfgNode): Configuration node + cfg_fname (string): Filename for the yaml format configuration file + name_tag (string): Additional name tag to identify this execution of the + configuration file, specified in :obj:`cfg.name_tag` + """ + run_name = os.path.splitext(os.path.basename(cfg_fname))[0] + run_name += f"-{name_tag}" if name_tag else "" + cfg.out_dir = os.path.join(cfg.out_dir, run_name) + + +def custom_set_run_dir(cfg, run_id): + """Custom output directory naming for each experiment run. + + Args: + cfg (CfgNode): Configuration node + run_id (int): Main for-loop iter id (the random seed or dataset split) + """ + cfg.run_dir = os.path.join(cfg.out_dir, str(run_id)) + # Make output directory + if cfg.train.auto_resume: + os.makedirs(cfg.run_dir, exist_ok=True) + else: + makedirs_rm_exist(cfg.run_dir) + + +def run_loop_settings(): + """Create main loop execution settings based on the current cfg. + + Configures the main execution loop to run in one of two modes: + 1. 'multi-seed' - Reproduces default behaviour of GraphGym when + args.repeats controls how many times the experiment run is repeated. + Each iteration is executed with a random seed set to an increment from + the previous one, starting at initial cfg.seed. + 2. 'multi-split' - Executes the experiment run over multiple dataset splits, + these can be multiple CV splits or multiple standard splits. The random + seed is reset to the initial cfg.seed value for each run iteration. + + Returns: + List of run IDs for each loop iteration + List of rng seeds to loop over + List of dataset split indices to loop over + """ + if len(cfg.run_multiple_splits) == 0: + # 'multi-seed' run mode + num_iterations = args.repeat + seeds = [cfg.seed + x for x in range(num_iterations)] + split_indices = [cfg.dataset.split_index] * num_iterations + run_ids = seeds + else: + # 'multi-split' run mode + if args.repeat != 1: + raise NotImplementedError("Running multiple repeats of multiple " + "splits in one run is not supported.") + num_iterations = len(cfg.run_multiple_splits) + seeds = [cfg.seed] * num_iterations + split_indices = cfg.run_multiple_splits + run_ids = split_indices + return run_ids, seeds, split_indices + + +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) + custom_set_out_dir(cfg, args.cfg_file, cfg.name_tag) + dump_cfg(cfg) + # Set Pytorch environment + torch.set_num_threads(cfg.num_threads) + # Repeat for multiple experiment runs + for run_id, seed, split_index in zip(*run_loop_settings()): + # Set configurations for each run + custom_set_run_dir(cfg, run_id) + set_printing() + cfg.dataset.split_index = split_index + cfg.seed = seed + cfg.run_id = run_id + seed_everything(cfg.seed) + auto_select_device() + if cfg.pretrained.dir: + cfg = load_pretrained_model_cfg(cfg) + logging.info(f"[*] Run ID {run_id}: seed={cfg.seed}, " + f"split_index={cfg.dataset.split_index}") + logging.info(f" Starting now: {datetime.datetime.now()}") + # Set machine learning pipeline + loaders = create_loader() + loggers = create_logger() + model = create_model() + if cfg.pretrained.dir: + model = init_model_from_pretrained( + model, cfg.pretrained.dir, cfg.pretrained.freeze_main, + cfg.pretrained.reset_prediction_head, seed=cfg.seed + ) + optimizer = create_optimizer(model.parameters(), + new_optimizer_config(cfg)) + scheduler = create_scheduler(optimizer, new_scheduler_config(cfg)) + # Print model info + logging.info(model) + logging.info(cfg) + cfg.params = params_count(model) + logging.info('Num parameters: %s', cfg.params) + # Start training + if cfg.train.mode == 'standard': + if cfg.wandb.use: + logging.warning("[W] WandB logging is not supported with the " + "default train.mode, set it to `custom`") + datamodule = GraphGymDataModule() + train(model, datamodule, logger=True) + else: + train_dict[cfg.train.mode](loggers, loaders, model, optimizer, + scheduler) + # Aggregate results from different seeds + try: + agg_runs(cfg.out_dir, cfg.metric_best) + except Exception as e: + logging.info(f"Failed when trying to aggregate multiple runs: {e}") + # When being launched in batch mode, mark a yaml as done + if args.mark_done: + os.rename(args.cfg_file, f'{args.cfg_file}_done') + logging.info(f"[*] All done: {datetime.datetime.now()}") From e6c2af82469892a6b66705f6228b099ba9eeaf0e Mon Sep 17 00:00:00 2001 From: Eugene29 <52263376+Eugene29@users.noreply.github.com> Date: Sun, 17 Dec 2023 16:57:27 -0800 Subject: [PATCH 4/5] Device agnostic --- graphgps/loader/master_loader.py | 3 +++ 1 file changed, 3 insertions(+) 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']] From 9aa289974b54e8ee0d2cbb6402c0e78b6f4e51e6 Mon Sep 17 00:00:00 2001 From: Eugene29 <52263376+Eugene29@users.noreply.github.com> Date: Sun, 17 Dec 2023 17:14:33 -0800 Subject: [PATCH 5/5] Device agnostic --- graphgps/transform/posenc_stats.py | 72 ++++++++++++++++++++---------- 1 file changed, 48 insertions(+), 24 deletions(-) diff --git a/graphgps/transform/posenc_stats.py b/graphgps/transform/posenc_stats.py index ff21452f..afe9ca53 100644 --- a/graphgps/transform/posenc_stats.py +++ b/graphgps/transform/posenc_stats.py @@ -8,7 +8,7 @@ to_undirected, to_dense_adj, scatter) from torch_geometric.utils.num_nodes import maybe_num_nodes from graphgps.encoder.graphormer_encoder import graphormer_pre_processing -import cupy as cp + def compute_posenc_stats(data, pe_types, is_undirected, cfg): """Precompute positional encodings for the given graph. @@ -52,14 +52,25 @@ def compute_posenc_stats(data, pe_types, is_undirected, cfg): # Eigen values and vectors. evals, evects = None, None - dev = torch.device("cuda") + 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: - undir_edge_index = undir_edge_index.to(dev) - 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) + 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 @@ -78,11 +89,20 @@ 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 - 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) + 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, @@ -104,17 +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: - ## is normalization None here on purpose? - 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) - -# 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)