TF-GNN训练图自编码器遇ValueError：无法解批标量GraphPiece

问题场景 #

尝试构建图自编码器实现图重构，训练阶段触发以下错误：

ValueError: Could not unbatch scalar (rank=0) GraphPiece.

相关代码如下：

创建GraphTensor #

edge_sources = train_data['from_encoded'].values
edge_targets = train_data['to_encoded'].values
edge_signs = train_data['sign'].values

node_ids = tf.concat([edge_sources, edge_targets], axis=0)
unique_ids, _ = tf.unique(node_ids)

num_nodes = len(unique_ids)
hidden_state_dim = 64

# Create a GraphTensor from edges and node features
graph = tf_gnn.GraphTensor.from_pieces(
    node_sets={
        "nodes": tf_gnn.NodeSet.from_fields(sizes=[num_nodes], features={'id': unique_ids, 'hidden_state': tf.zeros((num_nodes, hidden_state_dim))})
    },
    edge_sets={
        "edges": tf_gnn.EdgeSet.from_fields(
            sizes=[len(edge_sources)],
            adjacency=tf_gnn.Adjacency.from_indices(
                source=("nodes", edge_sources),
                target=("nodes", edge_targets)
            ),
            features={
                "sign": tf.convert_to_tensor(edge_signs, dtype=tf.float32)
            }
        )
    }
)

定义模型 #

graph_tensor_spec = graph.spec

# Define the GCN model with specified hidden layers
gcn_model = gcn.GCNConv(
        units=64,  # Example hidden layer sizes
        activation='relu',
        use_bias=True
    )
    
# Input layer using the graph tensor spec
inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)

# Apply the GCN model to the inputs
graph_setup = gcn_model(inputs,  edge_set_name="edges")
    
# Extract node states
node_states = graph_setup

decoder = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(64, activation='sigmoid')
])

decoded = decoder(node_states)

autoencoder = tf.keras.Model(inputs=inputs, outputs=decoded)

训练设置 #

autoencoder.compile(optimizer='adam', loss=tf.keras.losses.BinaryCrossentropy(),
    metrics=[tf.keras.metrics.AUC()])
autoencoder.fit(
    x=graph,
    y=graph,  # For autoencoders, input = output
    epochs=1   # Number of training epochs
)

错误原因 #

TensorFlow GNN的fit方法要求输入为批量GraphTensor（rank=1），但当前代码创建的是单个图的标量GraphTensor（rank=0），框架无法对其执行unbatch（拆分成批量元素）操作，因此报错。同时，直接将整个GraphTensor作为训练目标y也不符合自编码器的训练逻辑，需要明确重构的目标特征（如节点的原始隐藏状态）。

解决方案 #

1. 将单个GraphTensor转为批量格式 #

使用tf_gnn.batch或tf.expand_dims将单个图包装成批量形式：

# 方法1：用tf_gnn.batch包装
batched_graph = tf_gnn.batch([graph])

# 方法2：用tf.expand_dims扩展维度
batched_graph = tf.expand_dims(graph, 0)

2. 修正模型的节点特征提取 #

GCNConv的输出是更新后的GraphTensor，需要从中提取节点的hidden_state特征再传入解码器：

# 修正模型部分代码
graph_setup = gcn_model(inputs, edge_set_name="edges")
# 从GraphTensor中提取节点的hidden_state
node_states = graph_setup.node_sets["nodes"]["hidden_state"]
decoded = decoder(node_states)

3. 明确训练目标并批量处理 #

自编码器的目标是重构节点的原始特征，因此需要提取原始节点特征作为y，并转为批量格式：

# 提取原始节点隐藏状态作为训练目标
target_features = graph.node_sets["nodes"]["hidden_state"]
# 转为批量格式
batched_target = tf.expand_dims(target_features, 0)

4. 调整训练代码 #

使用批量后的GraphTensor和目标特征执行训练：

autoencoder.compile(optimizer='adam', loss=tf.keras.losses.BinaryCrossentropy(),
    metrics=[tf.keras.metrics.AUC()])
autoencoder.fit(
    x=batched_graph,
    y=batched_target,
    epochs=1
)

完整修正后的代码示例 #

# 1. 创建GraphTensor（原代码不变）
edge_sources = train_data['from_encoded'].values
edge_targets = train_data['to_encoded'].values
edge_signs = train_data['sign'].values

node_ids = tf.concat([edge_sources, edge_targets], axis=0)
unique_ids, _ = tf.unique(node_ids)

num_nodes = len(unique_ids)
hidden_state_dim = 64

graph = tf_gnn.GraphTensor.from_pieces(
    node_sets={
        "nodes": tf_gnn.NodeSet.from_fields(sizes=[num_nodes], features={'id': unique_ids, 'hidden_state': tf.zeros((num_nodes, hidden_state_dim))})
    },
    edge_sets={
        "edges": tf_gnn.EdgeSet.from_fields(
            sizes=[len(edge_sources)],
            adjacency=tf_gnn.Adjacency.from_indices(
                source=("nodes", edge_sources),
                target=("nodes", edge_targets)
            ),
            features={
                "sign": tf.convert_to_tensor(edge_signs, dtype=tf.float32)
            }
        )
    }
)

# 2. 转为批量格式
batched_graph = tf_gnn.batch([graph])
target_features = graph.node_sets["nodes"]["hidden_state"]
batched_target = tf.expand_dims(target_features, 0)

# 3. 定义模型（修正节点特征提取）
graph_tensor_spec = batched_graph.spec

gcn_model = gcn.GCNConv(
        units=64,
        activation='relu',
        use_bias=True
    )

inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
graph_setup = gcn_model(inputs, edge_set_name="edges")
node_states = graph_setup.node_sets["nodes"]["hidden_state"]

decoder = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(64, activation='sigmoid')
])

decoded = decoder(node_states)
autoencoder = tf.keras.Model(inputs=inputs, outputs=decoded)

# 4. 训练
autoencoder.compile(optimizer='adam', loss=tf.keras.losses.BinaryCrossentropy(),
    metrics=[tf.keras.metrics.AUC()])
autoencoder.fit(
    x=batched_graph,
    y=batched_target,
    epochs=1
)

内容的提问来源于stack exchange，提问作者Coder1934