Sample-based gnn training method, device, equipment, medium and program product

Abstract

The application particularly relates to a GNN training method and device based on sampling, equipment, a medium and a program product, comprising the following steps: determining a divided subgraph of each computing device; pre-sampling the divided subgraph on a preset server respectively, obtaining a sampling frequency of each vertex according to a sampling result, and then obtaining a feature cache file; mapping the feature cache file to a target cache space, extracting a preset number of vertex features based on the mapping to obtain a local feature cache, and then obtaining vertex features corresponding to the divided subgraph; calculating a sub-gradient of each computing device based on the vertex features, aggregating the sub-gradient based on a preset gradient aggregation function to obtain an aggregated gradient, and sending the aggregated gradient to each computing device, so that each computing device re-executes the step of obtaining the vertex features corresponding to the divided subgraph of each computing device until a preset ending condition is met. Therefore, the problems of high delay and the like in the distributed system training GNN caused by the communication bottleneck are solved.

Description

Technical Field

[0001] The present application relates to the field of computer software and theoretical technology, and in particular to a sampling-based GNN training method, apparatus, device, medium and program product. Background Art

[0002] Graph Neural Networks (GNNs) are the latest advancement in the field of deep learning. They extend the capabilities of deep learning to non-Euclidean domains and can be used to capture the complex patterns and dependencies inherent in graphs. They have been successfully applied in various scenarios.

[0003] In the related art, a server and a trainer are run on each machine, where the trainer is responsible for training on the local machine, and the server is responsible for storing all the data required for training on the local machine; in data parallel training, each machine must save all the data it uses locally.

[0004] However, many real-world graphs contain millions of vertices and edges, and the feature vectors of the vertices may have high dimensions; in this case, it becomes difficult to load the entire graph topology and all vertex feature vectors onto the computing device, which is a huge challenge for resource-limited edge devices and needs to be solved urgently. Summary of the Invention

[0005] The present application provides a sampling-based GNN training method, apparatus, device, medium, and program product to solve problems such as high latency when training GNN in a distributed system due to communication bottlenecks.

[0006] The first embodiment of the present application provides a sampling-based GNN training method, comprising the following steps:

[0007] Determine a partition subgraph for each computing device;

[0008] Pre-sampling the partitioned subgraphs of each computing device on a preset server, obtaining a sampling frequency of each vertex according to the sampling results, and obtaining a feature cache file according to the sampling frequency of each vertex;

[0009] Mapping the feature cache file to a target cache space, extracting a preset number of vertex features based on the feature cache file to obtain a local feature cache, and obtaining vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache;

[0010] A GNN model calculation is performed based on the vertex features corresponding to the partitioned subgraph of each computing device to obtain the sub-gradient of each computing device, and the sub-gradient of each computing device is aggregated based on a preset gradient aggregation function to obtain an aggregated gradient, and the aggregated gradient is sent to each computing device, so that each computing device re-executes the step of obtaining the vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache based on the aggregated gradient until a preset end condition is met.

[0011] Optionally, determining a partition subgraph for each computing device includes:

[0012] calculating the actual computing capacity of each computing device;

[0013] According to the actual computing capability of each computing device, the original graph data is partitioned based on a preset graph partitioning algorithm to obtain a partitioned subgraph of each computing device.

[0014] Optionally, the actual computing capacity of each computing device is:

[0015] ;

[0016] in, is the actual computing capability of the i-th computing device, is the average time of the i-th computing device, is the maximum average time among the K selected computing devices.

[0017] Optionally, before mapping the feature cache file to the target cache space, the method further includes:

[0018] In the first step of the first training cycle, the preset Python library function is used to obtain the available memory space;

[0019] The available memory space is divided into a spare memory space and the target cache space, and a memory mapping file is created in the target cache space to map the feature cache file to the target cache space.

[0020] Optionally, obtaining the vertex features corresponding to the partitioned subgraphs of each computing device according to the local feature cache includes:

[0021] Controlling each computing device to sample the corresponding partitioned subgraph to obtain a sampling subgraph corresponding to each computing device;

[0022] Based on the vertex identifiers in the sampling subgraph corresponding to each computing device, the vertex features corresponding to the partitioned subgraph of each computing device are loaded from the local feature cache, and the vertex features not present in the local feature cache are loaded from the preset server.

[0023] The second embodiment of the present application provides a sampling-based GNN training device, including:

[0024] A determination module, configured to determine a partition subgraph for each computing device;

[0025] A sampling module, configured to pre-sample the partitioned subgraphs of each computing device on a preset server, obtain a sampling frequency of each vertex according to the sampling results, and obtain a feature cache file according to the sampling frequency of each vertex;

[0026] a mapping module, configured to map the feature cache file to a target cache space, extract a preset number of vertex features based on the feature cache file to obtain a local feature cache, and obtain vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache;

[0027] A training module is configured to perform GNN model calculation based on the vertex features corresponding to the partitioned subgraphs of each computing device, obtain the sub-gradients of each computing device, aggregate the sub-gradients of each computing device based on a preset gradient aggregation function to obtain aggregated gradients, and send the aggregated gradients to each computing device, so that each computing device re-executes the step of obtaining the vertex features corresponding to the partitioned subgraphs of each computing device based on the local feature cache based on the aggregated gradients until a preset end condition is met. Optionally,

[0028] Optionally, the determining module is specifically configured to:

[0029] calculating the actual computing capacity of each computing device;

[0030] According to the actual computing capability of each computing device, the original graph data is partitioned based on a preset graph partitioning algorithm to obtain a partitioned subgraph of each computing device.

[0031] Optionally, in the training module, the actual computing capacity of each computing device is:

[0032] ;

[0033] in, is the actual computing capability of the i-th computing device, is the average time of the i-th computing device, is the maximum average time among the K selected computing devices.

[0034] Optionally, the mapping module is further configured to:

[0035] In the first step of the first training cycle, the preset Python library function is used to obtain the available memory space;

[0036] The available memory space is divided into a spare memory space and the target cache space, and a memory mapping file is created in the target cache space to map the feature cache file to the target cache space.

[0037] Optionally, the mapping module is specifically configured to:

[0038] Controlling each computing device to sample the corresponding partitioned subgraph to obtain a sampling subgraph corresponding to each computing device;

[0039] Based on the vertex identifiers in the sampling subgraph corresponding to each computing device, the vertex features corresponding to the partitioned subgraph of each computing device are loaded from the local feature cache, and the vertex features not present in the local feature cache are loaded from the preset server.

[0040] The third aspect of the present application provides an electronic device, comprising: at least one processor; and a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor, and the instructions are configured to execute the sampling-based GNN training method as described in the above embodiment.

[0041] The fourth aspect of the present application provides a computer-readable storage medium on which a computer program is stored, which is executed by a processor to implement the sampling-based GNN training method as described in the above embodiment.

[0042] The fifth aspect of the present application provides a computer program product, which stores a computer program. When the program is executed by a processor, it implements the sampling-based GNN training method as described in the above embodiment.

[0043] Thus, after determining the partitioned subgraphs of each computing device, the partitioned subgraphs are pre-sampled on a preset server. The sampling frequency of each vertex is obtained based on the sampling results, and then a feature cache file is obtained. The feature cache file is mapped to the target cache space, and based on this, a preset number of vertex features are extracted to obtain a local feature cache, and then the vertex features corresponding to the partitioned subgraph are obtained. Based on this, the sub-gradients of each computing device are calculated, and the sub-gradients are aggregated based on a preset gradient aggregation function to obtain an aggregated gradient. The aggregated gradient is sent to each computing device, so that each computing device re-executes the steps of obtaining the vertex features corresponding to the partitioned subgraph of each computing device until the preset end condition is met. This solves the problem of high latency when training GNNs in distributed systems due to communication bottlenecks.

[0044] Additional aspects and advantages of the present application will be given in part in the description below, and in part will become apparent from the description below, or will be learned through practice of the present application. BRIEF DESCRIPTION OF THE DRAWINGS

[0045] The above and / or additional aspects and advantages of the present application will become apparent and easily understood from the following description of the embodiments in conjunction with the accompanying drawings, in which:

[0046] Figure 1 A schematic diagram of the sampling-based GNN training process in edge scenarios in related technologies;

[0047] Figure 2 A flowchart of a sampling-based GNN training method provided according to an embodiment of the present application;

[0048] Figure 3 A schematic diagram of the overall workflow of a sampling-based GNN training method provided according to one embodiment of the present application;

[0049] Figure 4 Schematic diagram of a block diagram of a sampling-based GNN training device according to an embodiment of the present application;

[0050] Figure 5 A schematic diagram of the structure of an electronic device provided according to an embodiment of the present application. DETAILED DESCRIPTION

[0051] The following describes in detail embodiments of the present application, examples of which are shown in the accompanying drawings, wherein the same or similar reference numerals throughout represent the same or similar elements or elements having the same or similar functions. The embodiments described below with reference to the accompanying drawings are exemplary and are intended to be used to explain the present application, and should not be construed as limiting the present application.

[0052] Before introducing the sampling-based GNN training method of the embodiment of the present application, a brief introduction to the GNN training process in the related art is given.

[0053] Specifically, if Figure 1 As shown, Figure 1 The following is a schematic diagram of the sampling-based GNN training process in the edge scenario in the related art. Assuming that there are N devices for training GNN, the sampling-based GNN training method in the edge scenario includes the following steps:

[0054] (1) The server performs graph sampling on the partitioned subgraph of each device: Before training, the original graph data is divided into N subgraphs using a graph partitioning algorithm, with each device corresponding to one subgraph. During training, sampling is performed on the subgraph corresponding to each device;

[0055] (2) Extract the corresponding vertex features based on the sampled subgraph: obtain the vertices contained in the sampled subgraph and extract the required vertex features from the complete vertex feature matrix;

[0056] (3) Loading the sampled subgraph and vertex features onto the device: The sampled subgraph and the corresponding vertex features extracted in the previous step are transferred from the server to the corresponding device;

[0057] (4) Perform model calculation on the device: Perform GNN model calculation based on the sampled subgraph and vertex features transmitted from the server.

[0058] Researchers have developed many distributed GNN model training systems, such as DistDGL. DistDGL runs a server and a trainer on each machine, where the trainer is responsible for training on the local machine and the server is responsible for storing all the data required for training on the local machine (graph structure and all vertex feature vectors). In other words, in data parallel training, each machine must save all the data it uses locally. However, in edge scenarios, the memory capacity of edge devices is very limited, and as the scale of the graph increases, it becomes difficult to load all the data required for training onto the local device. In short, systems like DistDGL do not need to consider the memory capacity of edge devices, and edge devices cannot directly use systems like DistDGL to train larger graphs.

[0059] Based on the above problems, this application proposes a sampling-based GNN training method, in which the partitioned subgraphs of each computing device are determined; the partitioned subgraphs are pre-sampled on the preset server, and the sampling frequency of each vertex is obtained according to the sampling results, and then the feature cache file is obtained; the feature cache file is mapped to the target cache space, and based on this, a preset number of vertex features are extracted to obtain a local feature cache, and then the vertex features corresponding to the partitioned subgraph are obtained; based on this, the sub-gradients of each computing device are calculated, and the sub-gradients are aggregated based on the preset gradient aggregation function to obtain the aggregated gradient, and the aggregated gradient is sent to each computing device, so that each computing device re-executes the steps of obtaining the vertex features corresponding to the partitioned subgraph of each computing device until the preset end condition is met. In this way, the problem of high latency in GNN training in distributed systems caused by communication bottlenecks is solved.

[0060] Specifically, Figure 2 A flowchart of a sampling-based GNN training method provided in an embodiment of the present application.

[0061] like Figure 2 As shown in Figure 2, the sampling-based GNN training method includes the following steps:

[0062] In step S201 , a partitioned subgraph for each computing device is determined.

[0063] Among them, subgraph division refers to the process of dividing the original graph data into multiple subgraphs according to the actual computing power of each computing device before training GNN, and assigning these subgraphs to different computing devices.

[0064] Specifically, the actual computing power of each computing device must be calculated. This is typically obtained experimentally. Devices are categorized by type and a representative device is randomly selected from each category. This device and the server are then used to train on a given dataset for several epochs, and the average training time is calculated. Finally, the maximum average training time across all representative devices is taken and divided by the average training time per device to determine the device's computing power. Devices with higher computing power have greater data processing capabilities. Next, the original graph data is partitioned using a pre-defined graph partitioning algorithm based on the actual computing power of each device. Specifically, subgraphs of varying sizes are allocated to each device based on its computing power, ensuring that data synchronization between devices is not idle due to differences in computing power. During the graph partitioning process, the neighboring vertices of redundantly retained vertices are considered. This reduces the need to obtain information from other devices during training, thereby reducing inter-device communication and improving training efficiency. Finally, the partitioned subgraphs are transferred to the corresponding computing devices for storage. This allows each device to sample from its locally stored subgraph during training, reducing inter-device data transmission.

[0065] Optionally, in some embodiments, the actual computing capacity of each computing device is:

[0066] ;

[0067] in, represents the computing capability of device i, represents the average time per cycle required for training using server and device i, with a total of K categories of devices.

[0068] It is understandable that by defining the actual computing power of each computing device as a ratio relative to the training cycle time, the differences in computing performance between different devices can be quantified and compared. This metric not only provides an important reference for load balancing, resource allocation, and performance tuning in distributed training, but also helps make more informed decisions when conducting cost-benefit analyses. Understanding the computing power of each device is key to achieving efficient distributed computing and optimization, ensuring that tasks can be more evenly distributed across devices, thereby maximizing resource utilization and improving overall training speed and efficiency. In actual applications, more factors may need to be considered to accurately assess the computing power of a device, but this basic metric still has important guiding significance.

[0069] Optionally, in some embodiments, determining the partitioned subgraph of each computing device includes: calculating the actual computing capacity of each computing device; and partitioning the original graph data based on a preset graph partitioning algorithm according to the actual computing capacity of each computing device to obtain the partitioned subgraph of each computing device.

[0070] It can be understood that by calculating the actual computing power of each computing device, computing tasks can be allocated more reasonably, ensuring that each device can handle a workload that matches its capabilities, which helps avoid resource waste or overload; graph partitioning based on device computing power can ensure that the size of the subgraphs processed by different devices adapts to their respective processing capabilities, making data synchronization between devices more efficient during training, reducing unnecessary waiting time and communication delays; when partitioning the graph, redundantly retaining neighbor vertices can reduce communication delays during sampling and gradient synchronization; the feature cache obtained by pre-sampling can reduce the communication delay of feature transmission during training. By recording the sampling frequency of the vertex and pre-loading the vertex features to the device according to the frequency, the frequency and total amount of feature transmission can be reduced.

[0071] In step S202, pre-sampling is performed on the partitioned subgraph of each computing device on a preset server, and a sampling frequency of each vertex is obtained according to the sampling result, and a feature cache file is obtained according to the sampling frequency of each vertex.

[0072] Among them, the feature cache file is a pre-created file that contains the most frequently accessed vertex features identified through the pre-sampling process before training.

[0073] Specifically, pre-sampling usually includes executing a sampling process for a certain period of time, recording the number of times a vertex is selected, including both the case of being selected as a sampling vertex and the case of being selected as a neighbor vertex; pre-sampling can help estimate the "hotness" of a vertex in future sampling processes, that is, the possibility of being selected; during the pre-sampling process, the number of times each vertex is selected is recorded, and these data will be used to evaluate the frequency of each vertex being sampled during training; specifically, a vertex may appear repeatedly in multiple samplings, sometimes as a sampling vertex and sometimes as a neighbor vertex, so it is important to count the total number of times a vertex is sampled, whether as a sampling vertex or as a neighbor vertex; based on the sampling frequency of the vertex, a feature cache file is created, and the vertices are usually sorted in descending order of sampling frequency, and the features of the most frequently sampled vertices are selected as part of the feature cache file. If the total vertex feature file size is less than 2GB, all features are saved.

[0074] In step S203, the feature cache file is mapped to the target cache space, and a preset number of vertex features are extracted based on the feature cache file to obtain a local feature cache, and vertex features corresponding to the partitioned subgraph of each computing device are obtained according to the local feature cache.

[0075] Among them, the target cache space is a memory area set up to accelerate data access during training, and its purpose is to store vertex features extracted from the feature cache file; the local feature cache is created on each computing device and is a mechanism for storing frequently accessed vertex features; vertex features refer to the information or attributes carried by each vertex during GNN training. This information will be used to update the vertex status or pass it to adjacent vertices during model training.

[0076] Specifically, the available memory space is divided into two parts, one part is used as a backup memory space, and the other part is used as a target cache space; a memory mapping file is created in the target cache space, which is used to map the feature cache file on the SD card to the memory. In this way, the feature cache data on the SD card can be directly accessed in the memory, thereby improving the access speed; based on the feature cache file mapped to the memory, a preset number of vertex features are extracted from it as a local feature cache; for example, if the size of the feature cache file on the SD card exceeds 2GB, the most frequently sampled vertex features can be selected for loading; according to the vertex identifiers in the partition subgraph corresponding to each computing device, the vertex features corresponding to the device partition subgraph are loaded from the local feature cache; if some vertex features are not found in the local cache, these vertex features are loaded from the preset server.

[0077] Optionally, in some embodiments, vertex features corresponding to the partitioned subgraph of each computing device are obtained based on the local feature cache, including: controlling each computing device to sample the corresponding partitioned subgraph to obtain the sampled subgraph corresponding to each computing device; based on the vertex identifiers in the sampled subgraph corresponding to each computing device, loading the vertex features corresponding to the partitioned subgraph of each computing device from the local feature cache, and loading the vertex features that do not exist in the local feature cache from a preset server.

[0078] It is understandable that this process not only optimizes resource utilization and reduces unnecessary data transmission and memory consumption, but also accelerates the training process and improves the accuracy and reliability of the model; it also reduces the cost of distributed training, enhances the system's fault tolerance and scalability, and provides strong support for processing large-scale graph data and complex machine learning tasks.

[0079] In step S204, the GNN model is calculated based on the vertex features corresponding to the partitioned subgraph of each computing device to obtain the sub-gradient of each computing device, and the sub-gradient of each computing device is aggregated based on the preset gradient aggregation function to obtain the aggregated gradient, and the aggregated gradient is sent to each computing device, so that each computing device re-executes the step of obtaining the vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache based on the aggregated gradient until the preset end condition is met.

[0080] Among them, sub-gradient refers to the local gradient calculated by each computing device based on the vertex features in the partitioned subgraph to which it is assigned; the preset gradient aggregation function is a method used to integrate the sub-gradients from all computing devices to obtain a global gradient; the preset end condition refers to the standard or indicator set during the training process to determine when to stop training. The preset end condition can be a threshold set in advance by the user, a threshold obtained through a limited number of experiments, or a threshold obtained through a limited number of computer simulations. No specific limitation is made here.

[0081] Specifically, in a distributed environment, the entire graph data is divided into different computing devices; each computing device is responsible for processing the subgraph assigned to it and calculating the features of the vertices in these subgraphs based on the local feature cache; each computing device calculates the sub-gradients of the corresponding subgraph based on the subgraph and vertex features it processes. These sub-gradients reflect the derivatives of the loss function for the subgraph under the current parameters; a preset gradient aggregation function is used to aggregate the sub-gradients of all computing devices to obtain a global aggregated gradient, which is usually achieved through a communication mechanism to ensure that the gradient information of all devices can be aggregated to a central node or exchanged through a distributed communication protocol; the aggregated gradient is sent back to each computing device to ensure that each device has the latest global gradient information for subsequent parameter updates; each computing device updates its locally stored model parameters based on the received aggregated gradient; then each device repeats the steps of dividing the subgraph and calculating the vertex features based on the updated parameters and possible new local feature cache; the above steps are repeated until the preset end condition is met.

[0082] It should be noted that the preset termination conditions include, but are not limited to, reaching a predetermined number of training rounds. A predetermined upper limit for training rounds is set, and training automatically stops when this predetermined number of rounds is reached. Loss function convergence: monitoring changes in the loss function during training. When the change in the loss function is less than a certain threshold or remains unchanged for several consecutive rounds, the model is considered to have converged and training can be stopped. Validation set performance: testing the model's performance on unseen data. Training can be terminated when performance on the validation set (such as accuracy, F1 score, etc.) reaches a predetermined target value or does not improve or even deteriorates within a certain number of training rounds. Early stopping: if, after a period of training, performance on the validation set deteriorates instead of improving, training is terminated. Rate of change: when the rate of change of a key metric during training falls below a preset threshold, it indicates that there is little room for model improvement and training can be considered terminated. By setting appropriate termination conditions, overtraining can be effectively avoided while ensuring that the model reaches optimal or satisfactory performance within a reasonable training time. These conditions can be used individually or in combination to suit different training scenarios and objectives.

[0083] Thus, the partitioned subgraph of each computing device is determined; the partitioned subgraph is pre-sampled on a preset server, and the sampling frequency of each vertex is obtained based on the sampling results, and then a feature cache file is obtained; the feature cache file is mapped to the target cache space, and based on this, a preset number of vertex features are extracted to obtain a local feature cache, and then the vertex features corresponding to the partitioned subgraph are obtained; based on this, the sub-gradient of each computing device is calculated, and the sub-gradient is aggregated based on a preset gradient aggregation function to obtain an aggregated gradient, and the aggregated gradient is sent to each computing device, so that each computing device re-executes the steps of obtaining the vertex features corresponding to the partitioned subgraph of each computing device until the preset end condition is met. This solves the problem of high latency when training GNNs in distributed systems due to communication bottlenecks.

[0084] Optionally, in some embodiments, before mapping the feature cache file to the target cache space, it also includes: in the first step of the first cycle of training, using a preset Python library function to obtain available memory space; dividing the available memory space into a spare memory space and a target cache space, and creating a memory mapping file in the target cache space to map the feature cache file to the target cache space.

[0085] As a result, data access speed is significantly improved and memory resource allocation is optimized, which not only reduces unnecessary memory allocation and disk I / O operations, but also enhances system stability and training efficiency. Through sophisticated memory management strategies, a solid foundation is laid for the training process, ensuring the smooth progress of subsequent steps.

[0086] The sampling-based GNN training method is described in detail below in conjunction with a sampling-based GNN training method according to an embodiment of the present application.

[0087] Specifically, if Figure 3 As shown, the sampling-based GNN training method has the following steps:

[0088] Step 1: Before training, the original graph data is partitioned on the server using the proposed communication-friendly graph partitioning algorithm. Based on the differences in computing power between devices, each computing device is partitioned into a corresponding subgraph of appropriate size. Neighboring vertices are redundantly retained within the subgraph. The subgraph is then transferred to the corresponding device for storage, allowing for local sampling during training. Traditional sampling-based GNN training systems fail to account for the high communication latency of edge scenarios when performing graph partitioning. They do not redundantly retain the neighboring vertices of a vertex, which requires frequent communication between devices during sampling. Furthermore, they do not account for differences in computing power between devices, instead uniformly partitioning subgraphs of the same size, resulting in significant idle waits when synchronizing data between devices. However, our graph partitioning algorithm addresses these issues.

[0089] In step 2, before training, the server performs a sampling process for two to three cycles based on the subgraphs partitioned in step 1. The probability of a vertex being sampled is recorded to estimate the "hotness" of the sampled vertex. Unlike other systems that use pre-sampling, ours records both sampled training vertices and their neighboring vertices, whereas these only record training vertices. Within a single cycle, some vertices may reappear in many mini-batches, sometimes as sampled vertices and sometimes as neighboring vertices. Our graph partitioning algorithm aims to reduce cross-device communication by retaining redundant neighbor vertices. It records not only how often a vertex is sampled, but also how often it is partitioned into a subgraph as a redundant neighbor. These two aspects are then combined to determine the "hotness" of the vertex. Because vertex features need to be loaded onto the device in both cases, our approach can reduce the communication overhead of feature transmission.

[0090] In step 3, during the GNN training phase, although each mini-batch requires a different set of vertices, the memory size occupied by each mini-batch is roughly the same. This is because each mini-batch consists of a sampled subgraph with the same number of seed vertices and their vertex features. Therefore, when training the first mini-batch of the first epoch, we do not use the feature cache and use Python library functions to obtain available memory capacity. We then reserve 100MB of spare memory and use the remaining space as cache space. Furthermore, although the Raspberry Pi's swap space memory capacity is only 2GB, it can be expanded by attaching an SD card. Therefore, before training, we perform pre-sampling on the server to obtain a 2GB feature cache file and save it to the Raspberry Pi's SD card. During training, after obtaining available cache space, we use numpy.memmap in Python to create a memory-mapped file, map the feature cache file on the SD card into memory, and extract the specified number of vertex features as cache.

[0091] Step 4. During training, at the beginning of each step in the training cycle, the device first samples the partitioned subgraph stored locally to obtain a sampled subgraph. Then, based on the vertex ID in the sampled subgraph, the corresponding vertex features are obtained. If the device has already cached all the vertex features of the subgraph locally, there is no need to request vertex features from the server; otherwise, the device first determines which vertex features in the sampled subgraph are in the local cache, loads these vertex features from the local cache, and then loads other vertex features from the server through the network. In addition, in order to further reduce the communication delay caused by feature transmission, a pipeline with overlapping communication and calculation is designed. When executing the model calculation of the current sampled subgraph, the next step of sampling and vertex feature transmission is performed at the same time.

[0092] Step 5: During training, each time a device completes the calculation of the current step and obtains the sub-gradient, it transmits its own sub-gradient to the server. After the server receives the sub-gradients of all devices, it aggregates the gradients according to the gradient aggregation function and then distributes the gradients to each device. After receiving the gradients, each device uses the model with updated parameters to perform the next calculation.

[0093] Furthermore, the sampling-based GNN training method is described in detail with reference to a specific embodiment.

[0094] The sampling-based GNN training method includes the following steps:

[0095] Step 1: Perform graph partitioning before training. If there are N edge devices (i.e. computing devices) and graph G is the original graph, execute the graph partitioning algorithm to divide graph G into N sub-graphs. 、 、…、 ;

[0096] Step 2: Pre-sampling is performed before training. Each subgraph partitioned in Step 1 is pre-sampled on the server, and the sampling frequency of each vertex is recorded. The overall result of the pre-sampling on the server is then truncated to obtain a 2GB feature cache file, which is then saved to the edge device.

[0097] Step 3: In the first step of the first cycle of training, feature cache is not used, and the Python library function is used to obtain the available memory capacity. Then, 100MB of spare memory is reserved and the rest is used as cache space. After obtaining the available cache space, numpy.memmap in Python is used to create a memory-mapped file, map the 2GB feature cache file into memory, and extract the specified number of vertex features as cache.

[0098] Step 4: During training, at the beginning of each step in the training cycle, the device first samples the locally stored partitioned subgraph to obtain a sampled subgraph. It then retrieves the corresponding vertex features based on the vertex IDs in the sampled subgraph, first accessing the local feature cache and only accessing data not in the local feature cache from the server.

[0099] Step 5: During training, after each step of the model calculation in the training cycle is completed, the device's sub-gradient 、 、…、 The server collects all sub-gradients and aggregates them using a gradient aggregation function. The aggregated gradient is then sent to each device. The device receives the new gradients and updates its parameters for the next step of model calculation. Furthermore, while the current step is being calculated, the next step of sampling and feature transmission is also performed simultaneously.

[0100] The embodiment of the present application utilizes a communication-friendly graph partitioning algorithm to reduce cross-device communication overhead during sampling, and at the same time adopts a feature caching method to reduce the communication delay of vertex feature transmission and improve device resource utilization.

[0101] According to the sampling-based GNN training method proposed in the embodiment of the present application, after determining the partitioned subgraph of each computing device, the partitioned subgraph is pre-sampled on the preset server, and the sampling frequency of each vertex is obtained based on the sampling result, and then the feature cache file is obtained; the feature cache file is mapped to the target cache space, and based on this, a preset number of vertex features are extracted to obtain a local feature cache, and then the vertex features corresponding to the partitioned subgraph are obtained; based on this, the sub-gradient of each computing device is calculated, and the sub-gradient is aggregated based on the preset gradient aggregation function to obtain the aggregated gradient, and the aggregated gradient is sent to each computing device, so that each computing device re-executes the step of obtaining the vertex features corresponding to the partitioned subgraph of each computing device until the preset end condition is met. In this way, problems such as high latency when training GNNs in distributed systems caused by communication bottlenecks are solved.

[0102] Next, a sampling-based GNN training device proposed according to an embodiment of the present application is described with reference to the accompanying drawings.

[0103] Figure 4 3 is a block diagram of a sampling-based GNN training device according to an embodiment of the present application.

[0104] like Figure 4 As shown, the sampling-based GNN training device 10 includes: a determination module 100, a sampling module 200, a mapping module 300 and a training module 400.

[0105] The determination module 100 is configured to determine a partition subgraph for each computing device;

[0106] The sampling module 200 is used to pre-sample the partitioned subgraphs of each computing device on a preset server, obtain the sampling frequency of each vertex according to the sampling results, and obtain a feature cache file according to the sampling frequency of each vertex;

[0107] A mapping module 300 is configured to map the feature cache file to a target cache space, extract a preset number of vertex features from the feature cache file to obtain a local feature cache, and obtain vertex features corresponding to the partitioned subgraph of each computing device based on the local feature cache;

[0108] The training module 400 is used to perform GNN model calculation based on the vertex features corresponding to the partitioned subgraph of each computing device, obtain the sub-gradient of each computing device, aggregate the sub-gradient of each computing device based on a preset gradient aggregation function to obtain the aggregated gradient, and send the aggregated gradient to each computing device, so that each computing device re-executes the step of obtaining the vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache based on the aggregated gradient until the preset end condition is met.

[0109] Optionally, the determination module 100 is specifically used to: calculate the actual computing capacity of each computing device; and partition the original graph data based on a preset graph partitioning algorithm according to the actual computing capacity of each computing device to obtain a partitioned subgraph for each computing device.

[0110] Optionally, in the training module 400, the actual computing capacity of each computing device is:

[0111] ;

[0112] in, is the actual computing capability of the i-th computing device, is the average time of the i-th computing device, is the maximum average time among the K selected computing devices.

[0113] Optionally, the mapping module 300 is also used to: in the first step of the first cycle of training, use a preset Python library function to obtain available memory space; divide the available memory space into a spare memory space and a target cache space, and create a memory mapping file in the target cache space to map the feature cache file to the target cache space.

[0114] Optionally, the mapping module 300 is specifically used to: control each computing device to sample the corresponding partitioned subgraph to obtain a sampling subgraph corresponding to each computing device; based on the vertex identifiers in the sampling subgraph corresponding to each computing device, load the vertex features corresponding to the partitioned subgraph of each computing device from the local feature cache, and load the vertex features that do not exist in the local feature cache from the preset server.

[0115] It should be noted that the above explanation of the sampling-based GNN training method embodiment is also applicable to the sampling-based GNN training device of this embodiment, and will not be repeated here.

[0116] According to the sampling-based GNN training device proposed in the embodiment of the present application, the partitioned subgraph of each computing device is determined; the partitioned subgraphs are pre-sampled on the preset server respectively, and the sampling frequency of each vertex is obtained according to the sampling results, and then the feature cache file is obtained; the feature cache file is mapped to the target cache space, and based on this, a preset number of vertex features are extracted to obtain a local feature cache, and then the vertex features corresponding to the partitioned subgraph are obtained; based on this, the sub-gradient of each computing device is calculated, and the sub-gradient is aggregated based on the preset gradient aggregation function to obtain the aggregated gradient, and the aggregated gradient is sent to each computing device, so that each computing device re-executes the steps of obtaining the vertex features corresponding to the partitioned subgraph of each computing device until the preset end condition is met. In this way, problems such as high latency when training GNNs in distributed systems caused by communication bottlenecks are solved.

[0117] Figure 5 This is a schematic diagram of the structure of an electronic device provided in an embodiment of the present application. The electronic device may include:

[0118] Memory 501 , processor 502 , and computer programs stored in the memory 501 and executable on the processor 502 .

[0119] When the processor 502 executes the program, the sampling-based GNN training method provided in the above embodiment is implemented.

[0120] Furthermore, the electronic device further includes:

[0121] The communication interface 503 is used for communication between the memory 501 and the processor 502 .

[0122] The memory 501 is used to store computer programs that can be run on the processor 502 .

[0123] The memory 501 may include a high-speed RAM memory, and may also include a non-volatile memory (non-volatile memory), such as at least one disk memory.

[0124] If the memory 501, processor 502, and communication interface 503 are implemented independently, the communication interface 503, memory 501, and processor 502 can be interconnected via a bus and communicate with each other. The bus can be an Industry Standard Architecture (ISA) bus, a Peripheral Component Interconnect (PCI) bus, or an Extended Industry Standard Architecture (EISA) bus. Buses can be divided into address buses, data buses, control buses, etc. For ease of representation, Figure 5 Only one thick line is used in the diagram, but this does not mean that there is only one bus or one type of bus.

[0125] Optionally, in a specific implementation, if the memory 501, the processor 502 and the communication interface 503 are integrated on a chip, the memory 501, the processor 502 and the communication interface 503 can communicate with each other through an internal interface.

[0126] The processor 502 may be a central processing unit (CPU), an application specific integrated circuit (ASIC), or one or more integrated circuits configured to implement the embodiments of the present application.

[0127] An embodiment of the present application also provides a computer-readable storage medium having a computer program stored thereon, characterized in that when the program is executed by a processor, the sampling-based GNN training method as described above is implemented.

[0128] An embodiment of the present application also provides a computer program product, which stores a computer program, and when the program is executed by a processor, implements the above sampling-based GNN training method.

[0129] In the description of this specification, the description with reference to the terms "one embodiment", "some embodiments", "example", "specific example", or "some examples" means that the specific features, structures, materials or characteristics described in conjunction with the embodiment or example are included in at least one embodiment or example of the present application. In this specification, the schematic representations of the above terms do not necessarily refer to the same embodiment or example. Moreover, the specific features, structures, materials or characteristics described can be combined in any one or N embodiments or examples in a suitable manner. In addition, those skilled in the art can combine and combine different embodiments or examples described in this specification and features of different embodiments or examples without contradiction.

[0130] Furthermore, the terms "first" and "second" are used for descriptive purposes only and should not be understood to indicate or imply relative importance or implicitly specify the number of technical features indicated. Thus, a feature specified as "first" or "second" may explicitly or implicitly include at least one such feature. In the description of this application, "N" means at least two, for example, two, three, etc., unless otherwise specifically defined.

[0131] Any process or method description in a flowchart or otherwise described herein may be understood to represent a module, fragment or portion of code comprising one or more executable instructions for implementing the steps of a custom logical function or process, and the scope of the preferred embodiments of the present application includes alternative implementations in which functions may be performed out of the order shown or discussed, including performing functions in a substantially simultaneous manner or in reverse order depending on the functions involved, which should be understood by those skilled in the art to which the embodiments of the present application belong.

[0132] It should be understood that various parts of this application can be implemented using hardware, software, firmware, or a combination thereof. In the above embodiment, the N steps or methods can be implemented using software or firmware stored in a memory and executed by a suitable instruction execution system. For example, if implemented using hardware, as in another embodiment, any one of the following technologies known in the art or a combination thereof can be used: a discrete logic circuit having logic gate circuits for implementing logical functions on data signals, an application-specific integrated circuit having suitable combinational logic gate circuits, a programmable gate array (PGA), a field-programmable gate array (FPGA), etc.

[0133] Those skilled in the art will understand that all or part of the steps in the method of the above embodiment can be completed by instructing related hardware through a program, and the program can be stored in a computer-readable storage medium. When the program is executed, it includes one or a combination of the steps of the method embodiment.

Claims

1. A sampling-based GNN training method, characterized in that: The following steps are involved: Determine a partition subgraph for each computing device; Pre-sampling the partitioned subgraphs of each computing device on a preset server, obtaining a sampling frequency of each vertex according to the sampling results, and obtaining a feature cache file according to the sampling frequency of each vertex, wherein obtaining the sampling frequency of each vertex according to the sampling results and obtaining the feature cache file according to the sampling frequency of each vertex includes: in the pre-sampling process, recording the total number of times each vertex is selected as a sampling vertex or a neighbor vertex, and evaluating the sampling frequency of each vertex during training according to the total number of times, and sorting all vertices in descending order based on the sampling frequency of each vertex, and selecting the features of the most frequently sampled vertex as part of the feature cache file; Mapping the feature cache file to a target cache space, extracting a preset number of vertex features based on the feature cache file to obtain a local feature cache, and obtaining vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache; Performing GNN model calculation based on vertex features corresponding to the partitioned subgraph of each computing device to obtain sub-gradients of each computing device, aggregating the sub-gradients of each computing device based on a preset gradient aggregation function to obtain aggregated gradients, and sending the aggregated gradients to each computing device, so that each computing device re-executes the step of obtaining vertex features corresponding to the partitioned subgraph of each computing device based on the local feature cache based on the aggregated gradients until a preset end condition is met; determining the partitioned subgraph of each computing device includes: calculating the actual computing capacity of each computing device; partitioning the original graph data based on a preset graph partitioning algorithm according to the actual computing capacity of each computing device to obtain the partitioned subgraph of each computing device; the actual computing capacity of each computing device is: ;in, For the The actual computing power of a computing device, is the average time per cycle required for training using the server and the i-th computing device, is the total number of categories of the computing devices; before mapping the feature cache file to the target cache space, it also includes: in the first step of the first cycle of training, using a preset Python library function to obtain available memory space; dividing the available memory space into a spare memory space and the target cache space, and creating a memory mapping file in the target cache space to map the feature cache file to the target cache space; obtaining the vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache, including: controlling each computing device to sample the corresponding partitioned subgraph to obtain the sampling subgraph corresponding to each computing device; based on the vertex identifiers in the sampling subgraph corresponding to each computing device, loading the vertex features corresponding to the partitioned subgraph of each computing device from the local feature cache, and loading the vertex features that do not exist in the local feature cache from the preset server.

2. A sampling-based GNN training device, characterized in that: include: A determination module, configured to determine a partition subgraph for each computing device; A sampling module, configured to pre-sample the partitioned subgraphs of each computing device on a preset server, obtain a sampling frequency of each vertex based on the sampling results, and obtain a feature cache file based on the sampling frequency of each vertex, wherein obtaining the sampling frequency of each vertex based on the sampling results and obtaining the feature cache file based on the sampling frequency of each vertex include: during the pre-sampling process, recording the total number of times each vertex is selected as a sampling vertex or a neighbor vertex, and evaluating the sampling frequency of each vertex during training based on the total number of times, and sorting all vertices in descending order based on the sampling frequency of each vertex, and selecting the features of the most frequently sampled vertex as part of the feature cache file; a mapping module, configured to map the feature cache file to a target cache space, extract a preset number of vertex features based on the feature cache file to obtain a local feature cache, and obtain vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache; A training module is configured to calculate the sub-gradients of each computing device based on the vertex features corresponding to the partitioned subgraphs of each computing device, aggregate the sub-gradients of each computing device based on a preset gradient aggregation function to obtain aggregated gradients, and send the aggregated gradients to each computing device, so that each computing device re-executes the step of obtaining the vertex features corresponding to the partitioned subgraphs of each computing device based on the local feature cache based on the aggregated gradients until a preset termination condition is met; determining the partitioned subgraphs of each computing device includes: calculating the actual computing capacity of each computing device; partitioning the original graph data based on the actual computing capacity of each computing device based on a preset graph partitioning algorithm to obtain the partitioned subgraphs of each computing device; the actual computing capacity of each computing device is: ;in, For the The actual computing power of a computing device, is the average time per cycle required for training using the server and the i-th computing device, is the total number of categories of the computing devices; before mapping the feature cache file to the target cache space, it also includes: in the first step of the first cycle of training, using a preset Python library function to obtain available memory space; dividing the available memory space into a spare memory space and the target cache space, and creating a memory mapping file in the target cache space to map the feature cache file to the target cache space; obtaining the vertex features corresponding to the partitioned subgraph of each computing device according to the local feature cache, including: controlling each computing device to sample the corresponding partitioned subgraph to obtain the sampling subgraph corresponding to each computing device; based on the vertex identifiers in the sampling subgraph corresponding to each computing device, loading the vertex features corresponding to the partitioned subgraph of each computing device from the local feature cache, and loading the vertex features that do not exist in the local feature cache from the preset server.

3. An electronic device, characterized in that: include: A memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the program to implement the sampling-based GNN training method according to claim 1.

4. A computer-readable storage medium having a computer program stored thereon, characterized in that: The program is executed by a processor to implement the sampling-based GNN training method as claimed in claim 1.

5. A computer program product comprising a computer program, characterized in that When the computer program is executed by a processor, the sampling-based GNN training method according to claim 1 is implemented.

Images