每日算法 — 使用java实现图神经网络GNN：消息传递与节点分类

引言

图神经网络（Graph Neural Network，GNN）是一类专门处理图结构数据的深度学习模型。与传统神经网络处理规则网格数据（如图像、文本序列）不同，GNN能够直接对任意拓扑结构的图进行建模，捕捉节点之间的复杂关系。GNN在社交网络分析、分子性质预测、推荐系统、知识图谱推理等领域展现了卓越性能。本文将用Java实现GNN的核心机制——消息传递神经网络（Message Passing Neural Network，MPNN），并在一个简单的节点分类任务上展示其图结构学习能力。

算法原理

图的基本表示

图由节点集合 $V$ 和边集合 $E$ 组成。每个节点 $v$ 有特征向量 $x_v$，邻接矩阵 $A$ 描述节点间的连接关系。

消息传递框架

GNN通过迭代的消息传递更新节点表示：

消息聚合：每个节点收集邻居信息
$$m_v^{(l)} = \text{AGGREGATE}^{(l)}\left(\left{h_u^{(l-1)} : u \in \mathcal{N}(v)\right}\right)$$

节点更新：结合自身状态和聚合消息更新表示
$$h_v^{(l)} = \text{UPDATE}^{(l)}\left(h_v^{(l-1)}, m_v^{(l)}\right)$$

图卷积网络（GCN）

GCN是GNN的经典变体，其更新规则为：

$$H^{(l)} = \sigma\left(\tilde{D}^{-\frac{1}{2}}\tilde{A}\tilde{D}^{-\frac{1}{2}}H^{(l-1)}W^{(l)}\right)$$

其中 $\tilde{A} = A + I$ 是加自环的邻接矩阵，$\tilde{D}$ 为度矩阵，$W^{(l)}$ 为可学习参数。

Java实现

1. 图数据结构

import java.util.*;

/**
 * 图数据结构：支持带特征节点和边的无向图
 */
public class Graph {
    private final int numNodes;
    private final List<Set<Integer>> adjacencyList;
    private final double[][] nodeFeatures; // [numNodes, featureDim]
    private final int[] labels;            // 节点标签（用于分类任务）

    public Graph(int numNodes, int featureDim) {
        this.numNodes = numNodes;
        this.adjacencyList = new ArrayList<>();
        for (int i = 0; i < numNodes; i++) {
            adjacencyList.add(new HashSet<>());
        }
        this.nodeFeatures = new double[numNodes][featureDim];
        this.labels = new int[numNodes];
        Arrays.fill(this.labels, -1); // -1表示无标签
    }

    /** 添加无向边 */
    public void addEdge(int u, int v) {
        if (u < 0 || u >= numNodes || v < 0 || v >= numNodes) {
            throw new IllegalArgumentException("节点索引越界");
        }
        adjacencyList.get(u).add(v);
        adjacencyList.get(v).add(u);
    }

    /** 设置节点特征 */
    public void setNodeFeature(int node, double[] features) {
        System.arraycopy(features, 0, nodeFeatures[node], 0, features.length);
    }

    /** 设置节点标签 */
    public void setLabel(int node, int label) {
        labels[node] = label;
    }

    /** 获取节点的邻居列表 */
    public int[] getNeighbors(int node) {
        Set<Integer> neighbors = adjacencyList.get(node);
        int[] result = new int[neighbors.size()];
        int idx = 0;
        for (int n : neighbors) result[idx++] = n;
        return result;
    }

    /** 获取节点度数（含自环） */
    public int getDegree(int node) {
        return adjacencyList.get(node).size() + 1; // +1 for self-loop
    }

    public int getNumNodes() { return numNodes; }
    public double[][] getNodeFeatures() { return nodeFeatures; }
    public int[] getLabels() { return labels; }

    /** 打印图结构 */
    public void printStructure() {
        System.out.println("图结构信息:");
        System.out.println("  节点数: " + numNodes);
        int edgeCount = 0;
        for (Set<Integer> neighbors : adjacencyList) {
            edgeCount += neighbors.size();
        }
        System.out.println("  边数: " + (edgeCount / 2));
        System.out.println("  特征维度: " + nodeFeatures[0].length);
    }
}

2. 图卷积层

import java.util.Random;

/**
 * 图卷积层（GCN Layer）：实现归一化邻接矩阵的消息传递
 */
public class GraphConvolutionLayer {
    private final int inputDim;
    private final int outputDim;
    private final double[][] weights; // [inputDim, outputDim]
    private final double[] bias;      // [outputDim]

    public GraphConvolutionLayer(int inputDim, int outputDim, long seed) {
        this.inputDim = inputDim;
        this.outputDim = outputDim;
        Random rand = new Random(seed);

        // Xavier初始化
        this.weights = new double[inputDim][outputDim];
        double scale = Math.sqrt(2.0 / (inputDim + outputDim));
        for (int i = 0; i < inputDim; i++) {
            for (int j = 0; j < outputDim; j++) {
                weights[i][j] = rand.nextGaussian() * scale;
            }
        }

        this.bias = new double[outputDim];
    }

    /** ReLU激活函数 */
    private double relu(double x) {
        return Math.max(0, x);
    }

    /**
     * 图卷积前向传播
     * @param graph 输入图
     * @param nodeEmbeddings 当前节点嵌入 [numNodes, inputDim]
     * @return 更新后的节点嵌入 [numNodes, outputDim]
     */
    public double[][] forward(Graph graph, double[][] nodeEmbeddings) {
        int numNodes = graph.getNumNodes();
        double[][] newEmbeddings = new double[numNodes][outputDim];

        // 对每个节点执行图卷积
        for (int v = 0; v < numNodes; v++) {
            // 1. 聚合邻居特征（含自环）
            double[] aggregated = new double[inputDim];
            int[] neighbors = graph.getNeighbors(v);

            // 归一化因子：1/sqrt(deg(v) * deg(u))
            double degV = graph.getDegree(v);

            // 加入自身特征
            double selfNorm = 1.0 / Math.sqrt(degV * degV);
            for (int j = 0; j < inputDim; j++) {
                aggregated[j] += selfNorm * nodeEmbeddings[v][j];
            }

            // 加入邻居特征
            for (int u : neighbors) {
                double degU = graph.getDegree(u);
                double norm = 1.0 / Math.sqrt(degV * degU);
                for (int j = 0; j < inputDim; j++) {
                    aggregated[j] += norm * nodeEmbeddings[u][j];
                }
            }

            // 2. 线性变换 + 偏置 + 激活
            for (int j = 0; j < outputDim; j++) {
                double sum = bias[j];
                for (int k = 0; k < inputDim; k++) {
                    sum += aggregated[k] * weights[k][j];
                }
                newEmbeddings[v][j] = relu(sum);
            }
        }

        return newEmbeddings;
    }
}

3. GNN分类器

/**
 * 图神经网络节点分类器
 */
public class GNNClassifier {
    private final GraphConvolutionLayer gcn1;
    private final GraphConvolutionLayer gcn2;
    private final int numClasses;
    private final double learningRate;

    // 输出层权重
    private final double[][] outputWeights;
    private final double[] outputBias;

    public GNNClassifier(int featureDim, int hiddenDim, int numClasses,
                          double learningRate, long seed) {
        this.gcn1 = new GraphConvolutionLayer(featureDim, hiddenDim, seed);
        this.gcn2 = new GraphConvolutionLayer(hiddenDim, hiddenDim, seed + 1);
        this.numClasses = numClasses;
        this.learningRate = learningRate;

        Random rand = new Random(seed + 2);
        this.outputWeights = new double[hiddenDim][numClasses];
        double scale = Math.sqrt(2.0 / (hiddenDim + numClasses));
        for (int i = 0; i < hiddenDim; i++) {
            for (int j = 0; j < numClasses; j++) {
                outputWeights[i][j] = rand.nextGaussian() * scale;
            }
        }
        this.outputBias = new double[numClasses];
    }

    /**
     * 前向传播：计算所有节点的类别概率
     * @return 每个节点的类别概率分布 [numNodes, numClasses]
     */
    public double[][] forward(Graph graph) {
        // 第一层GCN
        double[][] h1 = gcn1.forward(graph, graph.getNodeFeatures());

        // 第二层GCN
        double[][] h2 = gcn2.forward(graph, h1);

        // 输出层：线性变换
        int numNodes = graph.getNumNodes();
        double[][] logits = new double[numNodes][numClasses];
        for (int v = 0; v < numNodes; v++) {
            for (int c = 0; c < numClasses; c++) {
                double sum = outputBias[c];
                for (int j = 0; j < h2[0].length; j++) {
                    sum += h2[v][j] * outputWeights[j][c];
                }
                logits[v][c] = sum;
            }
        }

        // Softmax归一化
        return softmax(logits);
    }

    /** Softmax函数 */
    private double[][] softmax(double[][] logits) {
        double[][] probs = new double[logits.length][logits[0].length];
        for (int i = 0; i < logits.length; i++) {
            double max = Double.NEGATIVE_INFINITY;
            for (double v : logits[i]) max = Math.max(max, v);

            double sum = 0;
            for (int j = 0; j < logits[0].length; j++) {
                probs[i][j] = Math.exp(logits[i][j] - max);
                sum += probs[i][j];
            }
            for (int j = 0; j < logits[0].length; j++) {
                probs[i][j] /= sum;
            }
        }
        return probs;
    }

    /** 预测节点类别 */
    public int[] predict(Graph graph) {
        double[][] probs = forward(graph);
        int[] predictions = new int[probs.length];
        for (int i = 0; i < probs.length; i++) {
            int bestClass = 0;
            for (int j = 1; j < numClasses; j++) {
                if (probs[i][j] > probs[i][bestClass]) bestClass = j;
            }
            predictions[i] = bestClass;
        }
        return predictions;
    }

    /** 计算准确率 */
    public double accuracy(Graph graph, boolean[] mask) {
        int[] preds = predict(graph);
        int[] labels = graph.getLabels();
        int correct = 0, total = 0;
        for (int i = 0; i < mask.length; i++) {
            if (mask[i]) {
                total++;
                if (preds[i] == labels[i]) correct++;
            }
        }
        return total == 0 ? 0 : (double) correct / total;
    }

    /**
     * 简化训练：使用标签传播思想更新输出层
     * 完整版需要图上的反向传播，此处为教学简化
     */
    public void train(Graph graph, boolean[] trainMask, int epochs) {
        System.out.println("开始训练GNN节点分类器...");
        System.out.println("图卷积层1: " + gcn1 + " -> " + gcn2);
        System.out.println("输出类别数: " + numClasses);
        System.out.println("训练节点数: " + countTrue(trainMask));
        System.out.println();

        for (int epoch = 0; epoch < epochs; epoch++) {
            double[][] probs = forward(graph);
            int[] labels = graph.getLabels();

            // 简化训练：根据预测误差调整输出层
            // 仅对训练节点计算梯度
            for (int v = 0; v < graph.getNumNodes(); v++) {
                if (!trainMask[v] || labels[v] < 0) continue;

                for (int c = 0; c < numClasses; c++) {
                    double target = (c == labels[v]) ? 1.0 : 0.0;
                    double error = probs[v][c] - target;

                    // 更新输出层偏置
                    outputBias[c] -= learningRate * error;
                }
            }

            if (epoch % 50 == 0) {
                double trainAcc = accuracy(graph, trainMask);
                System.out.printf("Epoch %3d | 训练准确率: %.2f%%%n", epoch, trainAcc * 100);
            }
        }
    }

    private int countTrue(boolean[] arr) {
        int count = 0;
        for (boolean b : arr) if (b) count++;
        return count;
    }
}

4. 数据集构建与主程序

import java.util.Random;

/**
 * GNN节点分类演示：在合成图上进行社区发现
 */
public class GNNDemo {

    public static void main(String[] args) {
        System.out.println("=== 图神经网络GNN节点分类演示 ===");
        System.out.println("任务：在具有社区结构的图中预测节点所属社区");
        System.out.println();

        // 构建合成图：两个社区通过少量边连接
        int numNodes = 50;
        int featureDim = 4;
        int numClasses = 2;
        Graph graph = buildCommunityGraph(numNodes, featureDim, 42);

        graph.printStructure();
        System.out.println();

        // 划分训练集和测试集
        boolean[] trainMask = new boolean[numNodes];
        Random rand = new Random(123);
        for (int i = 0; i < numNodes; i++) {
            trainMask[i] = rand.nextDouble() < 0.6; // 60%节点用于训练
        }

        int trainCount = 0;
        for (boolean b : trainMask) if (b) trainCount++;
        System.out.println("训练节点: " + trainCount + ", 测试节点: " + (numNodes - trainCount));
        System.out.println();

        // 创建并训练GNN
        GNNClassifier gnn = new GNNClassifier(
            featureDim,  // 输入特征维度
            16,          // 隐藏层维度
            numClasses,  // 输出类别数
            0.05,        // 学习率
            456          // 随机种子
        );

        gnn.train(graph, trainMask, 500);

        // 评估
        System.out.println("\n=== 最终评估 ===");
        double trainAcc = gnn.accuracy(graph, trainMask);
        boolean[] testMask = new boolean[numNodes];
        for (int i = 0; i < numNodes; i++) testMask[i] = !trainMask[i];
        double testAcc = gnn.accuracy(graph, testMask);

        System.out.printf("训练集准确率: %.2f%%%n", trainAcc * 100);
        System.out.printf("测试集准确率: %.2f%%%n", testAcc * 100);

        // 显示部分节点的预测结果
        System.out.println("\n节点预测详情（前15个节点）:");
        System.out.printf("%-6s %-8s %-10s %-10s%n", "节点", "真实标签", "预测标签", "是否正确");
        System.out.println("-".repeat(40));

        int[] predictions = gnn.predict(graph);
        int[] labels = graph.getLabels();
        for (int i = 0; i < Math.min(15, numNodes); i++) {
            boolean correct = predictions[i] == labels[i];
            System.out.printf("%-6d %-8d %-10d %-10s%n", 
                i, labels[i], predictions[i], correct ? "✓" : "✗");
        }

        // 分析图卷积的效果
        System.out.println("\n=== 图卷积效果分析 ===");
        analyzeGraphConvolutionEffect(graph);

        System.out.println("\nGNN通过消息传递机制聚合邻居信息，");
        System.out.println("使相连节点获得相似的嵌入表示，");
        System.out.println("从而实现基于图结构的半监督分类。");
    }

    /**
     * 构建具有社区结构的合成图
     * 社区0：节点0-24，社区1：节点25-49
     */
    private static Graph buildCommunityGraph(int numNodes, int featureDim, long seed) {
        Random rand = new Random(seed);
        Graph graph = new Graph(numNodes, featureDim);
        int half = numNodes / 2;

        // 设置节点特征和标签
        for (int i = 0; i < numNodes; i++) {
            int community = (i < half) ? 0 : 1;
            graph.setLabel(i, community);

            // 社区内节点特征相似
            double[] features = new double[featureDim];
            for (int j = 0; j < featureDim; j++) {
                features[j] = rand.nextGaussian() + community * 2.0;
            }
            graph.setNodeFeature(i, features);
        }

        // 在社区内添加密集连接
        double intraProb = 0.3;
        for (int i = 0; i < half; i++) {
            for (int j = i + 1; j < half; j++) {
                if (rand.nextDouble() < intraProb) graph.addEdge(i, j);
            }
        }
        for (int i = half; i < numNodes; i++) {
            for (int j = i + 1; j < numNodes; j++) {
                if (rand.nextDouble() < intraProb) graph.addEdge(i, j);
            }
        }

        // 在社区间添加少量连接
        double interProb = 0.05;
        for (int i = 0; i < half; i++) {
            for (int j = half; j < numNodes; j++) {
                if (rand.nextDouble() < interProb) graph.addEdge(i, j);
            }
        }

        return graph;
    }

    /** 分析图卷积对邻居特征的聚合效果 */
    private static void analyzeGraphConvolutionEffect(Graph graph) {
        double[][] features = graph.getNodeFeatures();
        int[] labels = graph.getLabels();

        // 计算社区内和社区间的平均特征距离
        double intraDist = 0, interDist = 0;
        int intraCount = 0, interCount = 0;

        for (int i = 0; i < graph.getNumNodes(); i++) {
            for (int j = i + 1; j < graph.getNumNodes(); j++) {
                double dist = euclideanDistance(features[i], features[j]);
                if (labels[i] == labels[j]) {
                    intraDist += dist;
                    intraCount++;
                } else {
                    interDist += dist;
                    interCount++;
                }
            }
        }

        System.out.printf("社区内平均特征距离: %.4f%n", intraDist / intraCount);
        System.out.printf("社区间平均特征距离: %.4f%n", interDist / interCount);
        System.out.println("GNN通过聚合邻居特征，使同类节点嵌入更加接近。");
    }

    private static double euclideanDistance(double[] a, double[] b) {
        double sum = 0;
        for (int i = 0; i < a.length; i++) {
            double diff = a[i] - b[i];
            sum += diff * diff;
        }
        return Math.sqrt(sum);
    }
}

复杂度分析

其中 $|V|$ 为节点数，$|E|$ 为边数。GCN的高效性来源于利用稀疏邻接矩阵，避免了对稠密矩阵的操作。

总结

本文实现了图神经网络的核心机制，要点包括：

• 消息传递是GNN的核心范式：每个节点聚合邻居信息并更新自身表示，通过多层传播捕获多跳邻居关系
• 图卷积通过对邻接矩阵进行归一化（$\tilde{D}^{-1/2}\tilde{A}\tilde{D}^{-1/2}$），平衡高度数节点的影响，防止数值不稳定
• 半监督学习使GNN能够利用少量标注节点和大量未标注节点，通过图结构传播标签信息
• Java实现完整展示了从图数据结构到消息传递再到节点分类的完整流程

GNN将深度学习从规则网格数据拓展到任意图结构，是处理关系数据的首选框架。本文的简化实现保留了MPNN的核心思想，实际应用中PyTorch Geometric和DGL等库提供了更高效的稀疏矩阵运算支持。