基于tensorflow写一个图神经网络（GNN）预测的python代码

以下是一个基于TensorFlow的GraphSAGE模型的代码示例，用于节点分类任务：

import tensorflow as tf
from tensorflow.keras import layers

class GraphSAGE(tf.keras.Model):
    def __init__(self, n_features, n_classes, n_hidden_layers=2, n_hidden_units=16, agg_method='mean'):
        super(GraphSAGE, self).__init__()
        self.n_features = n_features
        self.n_classes = n_classes
        self.n_hidden_layers = n_hidden_layers
        self.n_hidden_units = n_hidden_units
        self.agg_method = agg_method

        self.dense1 = layers.Dense(n_hidden_units, activation='relu')
        self.dense2 = layers.Dense(n_classes)

        self.aggregator_layers = []
        for i in range(n_hidden_layers):
            self.aggregator_layers.append(layers.Dense(n_hidden_units, activation='relu'))

    def call(self, inputs):
        x, adj_matrix = inputs

        # Aggregation
        for i in range(self.n_hidden_layers):
            if self.agg_method == 'mean':
                x = tf.matmul(adj_matrix, x)
                x = tf.divide(x, tf.reduce_sum(adj_matrix, axis=1, keepdims=True) + 1)
            elif self.agg_method == 'max':
                x = tf.matmul(adj_matrix, x)
                x = tf.reduce_max(x, axis=1, keepdims=True)
            else:
                raise ValueError('Invalid aggregation method')

            x = self.aggregator_layers[i](x)

        # Readout
        x = tf.reduce_mean(x, axis=0)
        x = self.dense1(x)
        x = self.dense2(x)

        return x

在这个代码中，我们定义了一个 `GraphSAGE` 类，它继承自 TensorFlow 的 `Model` 类。在 `__init__` 方法中，我们定义了模型的各种参数和层。在 `call` 方法中，我们定义了模型的前向传播过程。

我们的输入是一个大小为 `(n_nodes, n_features)` 的特征矩阵 `x` 和一个大小为 `(n_nodes, n_nodes)` 的邻接矩阵 `adj_matrix`。在聚合层中，我们使用邻接矩阵来聚合每个节点的邻居特征。我们可以使用平均值或者最大值来聚合邻居特征。在读出层，我们将所有节点的聚合表示取平均值，并将其输入到一个全连接层中，最后输出分类结果。

需要注意的是，这个代码示例中的 GraphSAGE 模型只是 GNN 中的一种，而且还有许多其他的 GNN 模型。不同的 GNN 模型可能具有不同的聚合方式和读出方式。