GNN training method, device and equipment based on sampling, medium and program product

Abstract

The invention particularly relates to a sampling-based GNN training method and device, equipment, a medium and a program product. The method comprises the following steps: determining a division sub-graph of each computing equipment; pre-sampling the divided sub-graphs on a preset server, obtaining the 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 feature cache file to obtain a local feature cache, and then obtaining vertex features corresponding to the divided sub-graphs; and calculating a sub-gradient of each computing device based on the sub-gradient, aggregating the sub-gradients 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 sub-graphs of each computing device. The preset ending condition is met. Therefore, the problems of high delay and the like when the distributed system trains the GNN due to 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, device, equipment, medium and program product. Background Art

[0002] Graph Neural Networks (GNN) are the latest development 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 vertices may have high dimensions; in this case, it becomes difficult to load the entire graph topology and all vertex feature vectors onto computing devices, 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 aspect 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 result, 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 subgraphs of each computing device according to the local feature cache;

[0010] 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] Among them, α i is the actual computing capacity of the ith computing device, t i 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, use the preset Python library function 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] Control 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 aspect of the present application provides a sampling-based GNN training device, including:

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

[0025] A sampling module, used 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 result, and obtain a feature cache file according to the sampling frequency of each vertex;

[0026] A mapping module, used 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 used 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] Among them, α i is the actual computing capacity of the ith computing device, t i 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 used to:

[0035] In the first step of the first training cycle, use the preset Python library function 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 used to:

[0038] Control 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, and the program 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] Therefore, after determining the partitioned subgraph of each computing device, 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-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, the problem of high latency when training GNN in a distributed system caused by communication bottlenecks is solved.

[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 the 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 an embodiment of the present application;

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

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

[0051] Embodiments of the present application are described in detail below, and examples of the embodiments 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, the GNN training process in the related art is briefly introduced.

[0053] Specifically, 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. Each device corresponds to one subgraph. During training, sampling is performed on the subgraph corresponding to each device.

[0055] (2) Extracting 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). That is, 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 size 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 partition subgraph of each computing device is determined; the partition subgraph is 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 partition 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 partition subgraph of each computing device until the preset end condition is met. In this way, the problem of high latency when training GNN in a distributed system 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, the sampling-based GNN training method includes the following steps:

[0062] In step S201 , a partition 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 and assigning these subgraphs to different computing devices according to the actual computing capabilities of each computing device before training the GNN.

[0064] Specifically, the actual computing power of each computing device needs to be calculated first, which is usually obtained through experiments, that is, the devices are classified according to the device type, and a device is randomly selected as a representative in each category. Then, the device and the server are used to train on a given data set for several cycles and the average time is calculated; finally, the maximum average time of all representative devices is taken and divided by the average time of each device to obtain the computing power of the device. The higher the computing power of the device, the stronger its ability to process data. Next, according to the actual computing power of each computing device, the original graph data is divided using a preset graph partitioning algorithm; specifically, different sizes of subgraphs are allocated to the device according to its computing power to ensure that there is no idle waiting when data is synchronized between devices due to differences in computing power; in the process of graph partitioning, the neighbor vertices of redundantly retained vertices need to be considered; this is to reduce the need to obtain information from other devices during training, thereby reducing communication between devices and improving training efficiency. Finally, the divided subgraphs are transferred to the corresponding computing devices for storage; in this way, during training, each device can perform sampling operations on its locally stored subgraphs to reduce data transmission between devices.

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

[0066]

[0067] Among them, α i represents the computing power of i, t 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 of different devices can be quantified and compared; this indicator not only provides an important reference for load balancing, resource allocation, and performance tuning in distributed training, but also helps to make more informed decisions in cost-benefit analysis; understanding the computing power of each device is the key to achieving efficient distributed computing and optimization, which ensures that tasks can be more evenly distributed among devices, thereby maximizing resource utilization and improving overall training speed and efficiency. In actual applications, more factors may need to be considered to accurately evaluate the computing power of a device, but this basic indicator still has important guiding significance.

[0069] Optionally, in some embodiments, determining the partitioned subgraph for each computing device includes: calculating the actual computing power of each computing device; and partitioning the original graph data based on a preset graph partitioning algorithm according to the actual computing power of each computing device to obtain the partitioned subgraph for 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 to ensure 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 subgraphs processed by different devices adapts to their respective processing capabilities, making data synchronization between devices more efficient during training and 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 subgraphs 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 "heat" 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 state of the vertex 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, and the file 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 to improve 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 according to 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, 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 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 the preset end condition is met.

[0080] Among them, sub-gradient refers to the local gradient calculated by each computing device according to 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 finite number of experiments, or a threshold obtained through a finite 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; using a preset gradient aggregation function, the sub-gradients of all computing devices are aggregated 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 end conditions include but are not limited to reaching the predetermined training rounds, setting a predetermined upper limit for the training rounds, and when the training reaches this predetermined round, the training process automatically stops; loss function convergence, monitoring the change of the loss function during the training process, when the change of the loss function is less than a certain threshold or does not change much in several consecutive cycles, it is considered that the model has converged and the training can be stopped; validation set performance, to detect the performance of the model on unseen data, when the performance on the validation set (such as accuracy, F1 score, etc.) reaches the predetermined target value, or does not improve or even deteriorates within a certain number of training cycles, the training can be terminated; early stopping strategy, if after a period of training, the performance on the validation set does not improve but decreases, then stop training; change rate, when the change rate of a key indicator during the training process is lower than the preset threshold, it means that there is little room for improvement in the model, and the training can be considered to be terminated. By setting appropriate end conditions, overtraining can be effectively avoided, while ensuring that the model reaches the best or satisfactory performance level within a reasonable training time; these conditions can be used alone or in combination to adapt to different training scenarios and goals.

[0083] Thus, the partition subgraph of each computing device is determined; the partition subgraph is pre-sampled on the preset server, and the sampling frequency of each vertex is obtained according to 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 partition 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 partition subgraph of each computing device until the preset end condition is met. Thus, the problem of high latency when training GNN in a distributed system due to communication bottlenecks is solved.

[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, the data access speed is significantly improved and the configuration of memory resources is optimized, which not only reduces unnecessary memory allocation and disk I / O operations, but also enhances the stability of the system and the 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, 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 according to the proposed communication-friendly graph partitioning algorithm. According to the difference in computing power of different devices, each computing device is partitioned into a corresponding partition subgraph of appropriate size, and the neighbor vertices are redundantly retained in the partition subgraph. The partition subgraph is then transferred to the corresponding device for storage, so that sampling can be performed locally on the device during training. Traditional sampling-based GNN training systems do not take into account the high communication latency in edge scenarios when performing graph partitioning. They do not redundantly retain the neighbor vertices of vertices, which requires frequent communication between devices during sampling. In addition, they do not take into account the difference in computing power between devices, but uniformly divide the partition subgraphs of the same size, resulting in serious idle waiting when synchronizing data between devices. However, our graph partitioning algorithm solves these problems.

[0089] Step 2: Before training, the server performs 2 to 3 cycles of sampling according to the subgraphs divided in step 1, and records the probability of the vertex being sampled to estimate the "heat" of the sampled vertex. Unlike other systems that use pre-sampling, we record both the sampled training vertices and the neighbor vertices, while they only record the training vertices. In one cycle, some vertices may reappear in many mini-batches, sometimes as sampled vertices and sometimes as neighbor vertices. Our graph partitioning algorithm aims to reduce cross-device communication by retaining redundant neighbor vertices. It not only records the frequency of vertex sampling, but also the frequency of vertex being divided into subgraphs as redundant neighbor vertices. These two aspects are then combined to determine the "heat" of the vertex. Because from the perspective of vertex feature usage, in both cases, the vertex features need to be loaded onto the device, so our method can reduce the communication overhead of feature transmission.

[0090] Step 3, during the training phase of GNN, 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 cycle, we do not use the feature cache and use the Python library function to obtain the available memory capacity, then reserve 100MB of spare memory and use the rest as cache space. In addition, although the swap space memory capacity of Raspberry Pi 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 SD card of Raspberry Pi. During training, after obtaining the 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, according to the vertex id in the sampled subgraph, the corresponding vertex features are obtained. If the device has cached all the vertex features of the subgraph locally, there is no need to request the 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 calculation of the current step and obtains a 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 each device receives the gradient, it uses the model with updated parameters to perform the next calculation.

[0093] Furthermore, the sampling-based GNN training method is described in detail in conjunction with 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 G. 1 , G 2 , …, GN ;

[0096] Step 2: Perform pre-sampling before training. Pre-sample each subgraph obtained in step 1 on the server and record the sampling frequency of each vertex. Then, truncate the overall result obtained by pre-sampling on the server to obtain a 2GB feature cache file and save it to the edge device.

[0097] Step 3: In the first step of the first cycle of training, feature cache is not used, and Python library functions are used to obtain 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 partitioned subgraph stored locally to obtain a sampled subgraph. Then, according to the vertex ID in the sampled subgraph, it obtains the corresponding vertex features, first accessing the local feature cache, and only accessing the server for data that is not in the local feature cache.

[0099] Step 5: During training, after each step of the model calculation in the training cycle, the device's sub-gradient grad 1 ,grad 2 , …, grad N The server collects all the sub-gradients and obtains the aggregated gradient according to the gradient aggregation function. The aggregated gradient is then sent to each device. The device receives the new gradient to update the parameters for the next model calculation. In addition, when executing the calculation of the current step, the next step of sampling and feature transmission will be performed at the same time.

[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 the resource utilization of the device.

[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 subgraphs are pre-sampled on the preset server, and the sampling frequency of each vertex is obtained according to 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 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 when training GNN in a distributed system due to communication bottlenecks is 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 It 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] Wherein, the determination module 100 is used to determine the partition subgraph of 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 result, and obtain a feature cache file according to the sampling frequency of each vertex;

[0107] A mapping module 300 is used to map the feature cache file to the 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;

[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, and 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] Among them, α i is the actual computing capacity of the ith computing device, t i 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 a 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 partition subgraph of each computing device is determined; the partition subgraph is pre-sampled on the preset server respectively, and the sampling frequency of each vertex is obtained according to 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 partition 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 partition subgraph of each computing device until the preset end condition is met. In this way, the problem of high latency when training GNN in a distributed system due to communication bottlenecks is solved.

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

[0118] A memory 501 , a processor 502 , and a computer program 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 comprises:

[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 executed 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, the processor 502 and the communication interface 503 are implemented independently, the communication interface 503, the memory 501 and the processor 502 can be connected to each other through a bus and communicate with each other. The bus can be an Industry Standard Architecture (ISA) bus, a Peripheral Component (PCI) bus or an Extended Industry Standard Architecture (EISA) bus. The bus can be divided into an address bus, a data bus, a control bus, 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 only 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), or 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, wherein when the program is executed by a processor, the above-mentioned sampling-based GNN training method 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" etc. 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 may be combined in any one or N embodiments or examples in a suitable manner. In addition, those skilled in the art may combine and combine the different embodiments or examples described in this specification and the features of the different embodiments or examples, without contradiction.

[0130] In addition, the terms "first" and "second" are used for descriptive purposes only and should not be understood as indicating or implying relative importance or implicitly indicating the number of technical features indicated. Therefore, a feature defined as "first" or "second" may explicitly or implicitly include at least one of the features. In the description of this application, "N" means at least two, such as two, three, etc., unless otherwise clearly and 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 not be performed in 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 technicians in the technical field to which the embodiments of the present application belong.

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

[0133] A person skilled in the art may understand that all or part of the steps in the method for implementing the above-mentioned embodiment may be completed by instructing related hardware through a program, and the program may be stored in a computer-readable storage medium, which, when executed, 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 result, and obtaining a feature cache file according to the sampling frequency of each vertex; 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 subgraphs of each computing device according to the local feature cache; 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.

2. The method according to claim 1, characterized in that The determining of the partition subgraph of each computing device includes: Calculating the actual computing capacity of each computing device; 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.

3. The method according to claim 2, characterized in that The actual computing capacity of each computing device is: Among them, α i is the actual computing capability of the ith computing device, t i is the average time of the i-th computing device, is the maximum average time among the K selected computing devices.

4. The method according to claim 1, characterized in that Before mapping the feature cache file to the target cache space, the method further includes: In the first step of the first training cycle, use the preset Python library function to obtain the available memory space; 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.

5. The method according to claim 1, characterized in that The obtaining, according to the local feature cache, vertex features corresponding to the partitioned subgraphs of each computing device includes: 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, 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.

6. A sampling-based GNN training device, characterized in that: The following steps are involved: A determination module, used to determine a partition subgraph for each computing device; A sampling module, used 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 result, and obtain a feature cache file according to the sampling frequency of each vertex; A mapping module, used 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 used to calculate the sub-gradient of each computing device based on the vertex features corresponding to the partitioned subgraph of each computing device, aggregate the sub-gradient of each computing device based on a preset gradient aggregation function to obtain an 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 a preset end condition is met.

7. The device according to claim 6, characterized in that The determination module is specifically used for: Calculating the actual computing capacity of each computing device; 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.

8. 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 as described in any one of claims 1 to 5.

9. 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 described in any one of claims 1-5.

10. 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 as described in any one of claims 1 to 5 is implemented.

Images