Graph neural network partitioning method based on heterogeneous GPU load balancing

Abstract

The invention discloses a graph neural network partitioning method based on heterogeneous GPU (Graphics Processing Unit) load balancing, relates to a GNN (Graphics Neural Network) technology, and provides a scheme for solving the problem that communication influence between partitions is not considered in the prior art. The method comprises the following steps: S1, acquiring an original image G and a partition target number N; s2, performing coarsening partitioning by using a graph partitioning algorithm; s3, randomly pre-distributing each sub-graph to each GPU (Graphics Processing Unit); s4, carrying out calculation overhead modeling on each GPU; s5, carrying out communication overhead modeling on each GPU; s6, setting a heuristic partition function and a memory constraint according to the calculation overhead and the communication overhead; s7, adjusting the coarsened partition according to the heuristic partition function and the memory constraint; and S8, distributing the partitioned sub-graphs to the corresponding GPUs, executing a graph calculation task, and summarizing and outputting the graph calculation task. The method has the advantages that the divided sub-graphs can not only adapt to the performance of the GPU, but also reduce communication requirements as much as possible and prevent memory overflow under the condition that the precision influence is small, and finally, the training time is effectively shortened, and the training efficiency is improved.

Images