A Cluster-Based Sampler for Fast GNN

Graph Neural Networks (GNNs) have become a powerful toolbox for solving complex problems dealing with graph-structured data. Research has gained significant momentum to improve the accuracy of GNN models. However, GNN training is computationally expensive. Several popular GNN methods use random-based neighborhood sampling techniques to reduce the GNN training time. These techniques aim to sparsify the graph by sampling the neighborhood vertices randomly, but they do not consider the graph structure and associated neighborhood features while sampling.We propose a cluster-based sampling technique that exploits the neighborhood structure and node features to address these challenges. To this, we first propose a weighted partitioning scheme that aims to partition the graph’s vertices by considering the similarity of vertex features and their neighborhood vertices. Then, we sample the neighborhood vertices based on the identified clusters and the common neighborhood vertices. We implemented the proposed sampling scheme in CUDA using the DGL framework. We evaluated our approach on large-scale datasets and observed that GNN, when used with our sampling scheme, shows a speedup of up to 2.6× compared to the GPU-based GraphSAGE implementation, which is a popular GNN based on the random sampling scheme.