Clustering rearrangement sampling method for accelerating graph neural network training
By clustering and edge rearrangement of nodes before training in graph neural networks, the problems of long sampling time and low memory utilization in the existing technology are solved, and more efficient GNN training is achieved, which improves cache hit rate and GPU memory utilization.
Patent Information
- Application Number
- CN202510567079.4
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-04-30
- Publication Date
- 2025-08-12
AI Technical Summary
In the training of existing graph neural networks, the sampling time is too long and the memory storage is discontinuous, resulting in the problems of too long training time and too low GPU memory utilization.
Before training, the nodes in the graph are clustered, the high-similarity nodes are divided into the same cluster, and the edge sets are rearranged according to the out and inlets of the connected nodes. The storage format adopts the COO format to improve the node storage efficiency and build a more compact subgraph to improve the GPU memory space utilization.
Through clustering and edge rearrangement, the cache hit rate and GPU bandwidth utilization efficiency are significantly improved, the time of each sampling iteration is reduced, and the efficiency and accuracy of GNN training are improved.
Smart Images
Figure FT_1 
Figure FT_2 
Figure FT_3
Abstract
Description
Technical Field
[0001] The present invention mainly relates to the application field of graph neural networks in deep learning, and in particular to a clustering and rearrangement sampling method for accelerating graph neural network training. Background Art
[0002] Graph Neural Networks (GNNs) are a new class of deep neural networks (DNNs) designed to process graph-structured data. They have attracted widespread attention from industry and academia due to their outstanding ability to learn and represent non-Euclidean graph data. Many GNNs have been widely used in various graph-related tasks, such as social network analysis, autonomous driving, and recommendation systems. In practice, many real-world graphs are large-scale and associated with rich node attributes (i.e., features). The goal of GNNs is to learn a low-dimensional feature representation (i.e., embedding) for each vertex in the graph and use this representation for graph-related tasks. Due to their high memory bandwidth and huge parallelism, GPUs have been widely used to accelerate GNN training. Graph learning systems such as PYG and DGL provide optimized GPU implementations and convenient programming interfaces to enhance the performance and usability of GNN training.
[0003] The training phase of a GNN typically centers on iterative information aggregation and feature updates, requiring repeated visits to neighboring nodes and vector calculations. However, in the real world, large-scale graph data exhibits typical characteristics such as sparseness, heterogeneity, and power-law distribution, leading to frequent random memory accesses and inefficient graphics memory utilization during training. To alleviate these bottlenecks, systems such as AliGraph proposed by Alibaba Labs, NeuGraph and PipeGCN proposed by the Microsoft Research Asia team, and GraphSAGE proposed by William L. Hamilton and others at McGill University have proposed system optimization strategies from the perspectives of caching mechanisms and pipeline execution.
[0004] Unlike existing research that is not based on clustered sampling, this study focuses on performance optimization in the GNN sampling phase and proposes a clustered reordering sampling method to accelerate graph neural network training, which is used to improve the sampling efficiency of graph neural networks. The core idea of this method is to use the clustering structure of nodes in the feature space in the graph to reorder the edge set, thereby enhancing the locality of the edges and enabling the GPU to more efficiently access continuously stored node features during training. This edge reordering strategy forms high-similarity node clusters by performing algorithmic clustering in the preprocessing phase before training, and reorders the edges based on the clustering results. This makes the memory access pattern more continuous and compact during edge sampling and message passing, significantly improving cache hit rate and GPU bandwidth utilization efficiency.
[0005] Current sampling in GNN training still suffers from issues such as long sampling times and discontinuous memory storage, leading to long training times and low GPU memory utilization. This is because sampling and computation often involve repeated access to nodes, and node locality is low. Using different sampling methods can lead to the loss of important nodes and features in the graph, resulting in reduced GNN training accuracy. Different sampling methods may also have various limitations. Summary of the Invention
[0006] The purpose of the present invention is to propose a clustering and rearrangement sampling method for accelerating graph neural network training in response to the current situation where GNN sampling has not been pre-processed by clustering and rearrangement. The core idea is: before training, clustering is performed first to divide high-similarity nodes into the same cluster, and the edge set in the graph is rearranged in the cluster. The rearrangement order of the edges is determined by the out-degree and in-degree of the connected nodes. The storage format adopts the COO format, which greatly improves the storage efficiency of the nodes. Since the rearranged edge sampling subgraph has better structural compactness, the subgraph construction is more efficient, the GPU memory space utilization is improved, and the time of each sampling iteration is reduced, thereby accelerating the GNN training time.
[0007] The present invention is implemented as follows: a clustering and rearrangement sampling method for accelerating graph neural network training includes the following steps: Step 1: Store the original image edge sequence in COO format; Step 2: Based on the given parameters, perform algorithmic clustering on the nodes in the graph; Step 3: Rearrange the edges according to the clustering results; Step 4: Calculate the available GPU memory and determine the number of cache nodes; Step 5: Load the data into GPU memory; Step 6: Perform GNN subsequent sampling and calculation operations.
[0008] In the above scheme, step 1 includes the following steps:
[0009] Step 1.1: Extract the starting index array B_arr and the ending index array E_arr from the rearranged edge sequence and construct a sparse adjacency structure in COO format.
[0010] Step 1.2: If the graph is a weighted graph, the corresponding weights are extracted at the same time to form a triple COO structure; if it is an unweighted graph, only the start and end point indices are stored.
[0011] In the above scheme, step 2 sets the number of clusters to be re-arranged based on the size and complexity of the graph. At the same time, the cluster number of each node is recorded to form the cluster number index table inx_arr.
[0012] In the above scheme, step 3 includes the following steps:
[0013] Step 3.1: Traverse each edge in the original graph and obtain the numbers of its starting and ending clusters.
[0014] Step 3.2: Calculate the out-degree and in-degree information of the node and the reordering priority index of the edge.
[0015] Step 3.3: Reorder the edges based on the cluster ID and node degree, giving priority to edges in the same cluster, and placing high-priority edges (high total degree) first in storage, thus generating a reordered edge sequence.
[0016] In the above scheme, step five loads the data obtained in step three into the GPU for subsequent calculations.
[0017] The beneficial effects of the present invention are as follows: (1) The present invention effectively aggregates similar nodes into adjacent storage locations through clustering. Nodes in the same cluster are compactly distributed in memory, which significantly improves the locality of node access and reduces random access overhead. (2) The reordered edge structure in the present invention reduces cross-cluster access, lowers the frequency and amount of CPU-GPU data transmission, enables each round of message aggregation operations to be completed on the GPU, avoids frequent memory queries and data migration, and thus significantly reduces the time consumption of each training iteration. (3) The present invention can accelerate training for graphs of different sizes, and the data is stored in COO format, which is compatible with many GNN frameworks and has strong versatility. BRIEF DESCRIPTION OF THE DRAWINGS
[0018] Figure 1 This is the overall flow chart of clustering and rearrangement sampling for graph neural network training of the present invention;
[0019] Figure 2 It is an overall flow chart of the implementation of the present invention;
[0020] Figure 3 is an undirected graph and its corresponding compressed sparse column (COO) representation in the present invention;
[0021] Figure 4 It is the specific implementation of the algorithm in the present invention on the undirected graph and the rearranged B_arr and E_arr;
[0022] Figure 5 is the total degree of each node in the present invention, and A represents the characteristic matrix (formula);
[0023] Figure 6 is the calculation formula of the node ranking score in the present invention;
[0024] Figure 7 is the GPU memory available for calculation in the present invention (formula). DETAILED DESCRIPTION
[0025] In order to make the purpose, technical solutions and advantages of the present invention more clearly understood, the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.
[0026] Example: The nodes in the Reddit dataset represent Reddit posts, and the edges indicate that a connection is established between two posts if the same user comments on both posts. This dataset comes from the Reddit social media platform, which contains a large amount of user interaction information. The features of the nodes are high-dimensional sparse vectors, which represent the behavioral characteristics of users on the platform (such as posting, commenting, liking, etc.), and these features are collected through user data on the platform. The label of each node indicates that the user belongs to a certain community or group, and these labels are clustered or manually labeled. The Reddit dataset contains 232,965 nodes, 11,606,909 edges, and the average degree of the nodes is 50, which means that on average each post is connected to 50 other posts through comments from the same user. Similar to the ogbn-proteins dataset, the graph structure and node features in the Reddit dataset are complex, and using traditional methods for calculation may encounter huge time and space overhead, especially when training on large-scale graph data. The steps for implementing a clustering and permutation sampling method to accelerate graph neural network training on the Reddit dataset are as follows:
[0027] Step 1: Store the edge sequence in COO format; Record the start and end indices of each edge in the edge sequence in the graph and store them in B_arr and E_arr respectively. Next, convert the two arrays into COO format.
[0028] Step 2: Based on the given parameters, perform algorithmic clustering on the nodes in the graph; Get the node feature information of Reddit and construct the corresponding node feature matrix. Sampling is performed on the overall data of the Reddit dataset, and our parameters are set as follows: the total number of clusters n_cluster = 3. The K-means algorithm is generally clustered based on feature vectors, but in this patent, it is clustered based on the total number of out-degree and in-degree of the node, and the same node will not appear in different clusters at the same time. Specifically, a node with a total degree of x is selected, and nodes with a node degree range of [x-1, x+1] will be clustered into this cluster, and a label will be set to indicate that the node has been selected. At the same time, the selected nodes will not be divided into other clusters, nor will they be selected as new cluster centers. Each cluster will generate a unique cluster ID, and the cluster IDs are incremented sequentially.
[0029] Clustering is performed according to the algorithm described above. Within the same cluster, we set the nodes to be arranged in descending order of degree (nodes with high degrees are given priority). At the same time, the cluster number of each node is recorded to form the cluster number index table inx_arr.
[0030] Step 3: Rearrange the edges according to the clustering results; From step 2, we obtain a cluster of successfully clustered nodes. To ensure high cache locality and continuous access to similar nodes, we need to rearrange the edges of the nodes. First, extract the source and target nodes of all edges from the original Reddit graph and stack them into COO format, which can also be converted into a matrix. Next, calculate the out-degree and in-degree of each node and add the two numbers to obtain the total degree, which is used as one of the sorting indicators. Another sorting indicator is the cluster ID. Check whether there is a pre-calculated cluster ID. If not, randomly assign an ID (0-2). Finally, sort in descending order according to the sorting formula (shown in Formula 2), and record the sorted B_arr and E_arr. The rearranged B_arr is recorded according to the total degree of the node. This is to prepare for insufficient GPU memory in the future and cache nodes with high degrees to ensure that training accuracy is not significantly affected.
[0031] Step 4: Calculate the available GPU memory and determine the number of cache nodes; After completing step 3, we need to load the data into the GPU for subsequent GNN training. Using the formula (shown in Equation 3), we calculate the available GPU memory and determine how much node data can be cached. If the available memory is large enough, all data can be cached; otherwise, only the top-ranked data can be cached.
[0032] Step 5: Load the data into GPU memory; After determining the amount of cacheable data in step 4, the feature vectors and node information of the rearranged edges and nodes can be stored on the GPU. This ensures that subsequent operations can directly use the GPU data, reducing data transfer time. If the GPU cannot store all the data, the remaining data is stored on the CPU.
[0033] Step 6: Perform GNN subsequent sampling and calculation operations. Through the preceding five steps, we have completed a clustered, re-ranked sampling method for accelerating graph neural network training, efficiently and accurately obtaining node data. GNN sampling and computation operations are performed using data already cached in the GPU. First, the cache is initialized based on memory availability, with important nodes already pre-loaded into GPU memory. During multi-layer sampling, nodes with cache hits are directly selected from the GPU to retrieve their corresponding node features; nodes with cache misses are loaded from the CPU. Neighbors are also retrieved based on the re-ranked indices, as similar nodes are clustered and located adjacent to each other in memory. Sampling is then performed using the previously calculated node degrees, batch sampling of clustered nodes to reduce random access time. In the feature aggregation phase, the contiguous storage of clustered similar nodes is exploited to parallelize processing and update frequently accessed nodes to GPU memory. In the message passing phase, computation is performed based on the re-ranked indices to maximize cache hit rate. This algorithmic re-ranked sampling significantly improves GNN training efficiency while maintaining model accuracy.
[0034] The process of the clustering and rearrangement sampling method proposed above to accelerate graph neural network training is as follows: Figure 1 As shown in the figure. Random access to node features in traditional GNNs results in extremely irregular memory access patterns, resulting in low cache hit rates and limited computational efficiency. Therefore, an algorithm-based node clustering and reordering mechanism is proposed. By grouping nodes with similar features and prioritizing highly ordered nodes, a more localized data access pattern is constructed. During training, the clustered node organization ensures that similar nodes are processed continuously, significantly improving the cache hit rate. This computational approach not only significantly reduces memory access latency but also reduces data transmission overhead between the GPU and CPU. This method can maintain model accuracy while also improving computational efficiency in the sampling phase, thereby improving overall system performance.
[0035] Although the above describes the specific embodiments of the present invention in conjunction with the accompanying drawings, it is not intended to limit the scope of protection of the present invention. Those skilled in the art should understand that various modifications or variations that can be made by those skilled in the art on the basis of the technical solution of the present invention without the need for creative work are still within the scope of protection of the present invention.
Claims
1. A clustering and re-sampling method for accelerating graph neural network training, characterized in that: The clustering and re-sampling method for accelerating graph neural network training includes the following steps: Step 1: Store the original image edge sequence in COO format; Step 2: Based on the given parameters, perform algorithmic clustering on the nodes in the graph; Step 3: Rearrange the edges according to the clustering results; Step 4: Calculate the available GPU memory and determine the number of cache nodes; Step 5: Load the data into GPU memory; Step 6: Perform GNN subsequent sampling and calculation operations.
2. The clustering and re-sampling method for accelerating graph neural network training according to claim 1, wherein: The step 1 stores the original image edge sequence in COO format. The specific steps are as follows: Step 1: Extract the starting index array B_arr and the end index array E_arr from the rearranged edge sequence and construct a sparse adjacency structure in COO format; Step 2: If the graph is a weighted graph, the corresponding weights are extracted at the same time to form a triple COO structure; if it is an unweighted graph, only the start and end point indexes are stored.
3. The clustering and re-sampling method for accelerating graph neural network training according to claim 1, wherein: The second step is to perform algorithmic clustering on the nodes in the graph according to given parameters; set the number of clustering rearrangements according to the size and complexity of the graph; and record the cluster number to which each node belongs to form a cluster number index table inx_arr.
4. The clustering and re-sampling method for accelerating graph neural network training according to claim 1, wherein: The step three is to rearrange the edges according to the obtained clustering results, and the specific steps are as follows: Step 1: Traverse each edge in the original graph and obtain the numbers of its starting and ending clusters; Step 2: Calculate the out-degree and in-degree information of the node and the reordering priority index of the edge; Step 3: Reorder the edges according to the cluster ID and node degree, so that the edges in the same cluster are arranged first, and the high-priority edges (high total degree) are placed first in the storage, thereby generating a reordered edge sequence.
5. The clustering and re-sampling method for accelerating graph neural network training according to claim 1, wherein: The fourth step is to calculate the available GPU memory and determine the number of cache nodes.
6. The clustering and re-sampling method for accelerating graph neural network training according to claim 1, wherein: The step five loads the data in step three into the GPU memory.
7. The clustering and re-sampling method for accelerating graph neural network training according to claim 1, wherein: Step 6 performs GNN subsequent sampling and calculation operations.