如何向tf.Estimator输入可变形状批量二维数组？图数据场景问询

我完全理解你的困扰——图数据的节点数天生不固定，而tf.estimator.inputs.numpy_input_fn确实只支持固定形状的批量数据，没法直接应对这种场景。不过用tf.data.Dataset构建输入管道就能轻松解决，下面是具体的实现步骤和示例：

核心思路 #

tf.data.Dataset支持生成可变形状的数据，搭配from_generator或者直接构造数据集，再通过填充（padding）或RaggedTensor来处理批量内的形状差异，最后封装成tf.Estimator能识别的输入函数。

步骤1：准备可变形状的图数据 #

假设你的每个图是一个二维数组（形状为[节点数N_i, 特征维度D]），标签是单个值，示例数据如下：

import numpy as np
import tensorflow as tf

# 3个示例图，节点数分别为3、5、4，特征维度为2
graph_features = [
    np.random.rand(3, 2).astype(np.float32),
    np.random.rand(5, 2).astype(np.float32),
    np.random.rand(4, 2).astype(np.float32)
]
graph_labels = np.array([0, 1, 0], dtype=np.int32)  # 每个图对应一个标签

步骤2：用tf.data.Dataset构建可变形状数据集 #

这里用from_generator来生成数据集，它能完美适配可变长度的数据：

def input_generator():
    """生成器函数，逐个返回图特征和标签"""
    for feat, label in zip(graph_features, graph_labels):
        yield feat, label

# 定义输出签名，用None标记可变维度（节点数）
dataset = tf.data.Dataset.from_generator(
    input_generator,
    output_signature=(
        tf.TensorSpec(shape=(None, 2), dtype=tf.float32),  # 可变节点数，固定特征维度
        tf.TensorSpec(shape=(), dtype=tf.int32)            # 标签是标量
    )
)

步骤3：处理批量数据（两种方案） #

方案A：填充到统一形状（适合传统模型）

如果你的模型需要固定形状的输入，可以用padded_batch把每个批量内的图填充到该批次的最大节点数：

# 打乱数据后按批次填充，batch_size设为你的需求
dataset = dataset.shuffle(buffer_size=len(graph_features)).padded_batch(
    batch_size=2,
    padded_shapes=((None, 2), ()),  # 仅对特征的节点维度填充，标签无需填充
    padding_values=(0.0, 0)         # 填充值可自定义，默认是0
)

方案B：使用RaggedTensor（适合图模型）

如果想保留原始节点数，避免填充带来的冗余，可以直接用batch生成RaggedTensor（TensorFlow 2.0+支持）：

dataset = dataset.shuffle(buffer_size=len(graph_features)).batch(2)
# 此时批量特征是RaggedTensor，形状为[batch_size, None, 2]

步骤4：封装成Estimator可用的输入函数 #

只需要把数据集转换成tf.Estimator要求的输入格式即可：

def my_input_fn():
    """供Estimator调用的输入函数"""
    # 如果是分类任务，需要把特征包装成字典（key对应模型的输入层名称）
    def map_func(features, labels):
        return {"graph_input": features}, labels
    
    return dataset.map(map_func)

步骤5：用Estimator训练模型 #

现在就可以像使用numpy_input_fn一样用这个输入函数了：

# 示例：定义一个简单的DNN分类器（如果用RaggedTensor，需要在模型里处理）
estimator = tf.estimator.DNNClassifier(
    feature_columns=[tf.feature_column.numeric_column("graph_input", shape=(None, 2))],
    hidden_units=[64, 32],
    n_classes=2
)

# 开始训练
estimator.train(input_fn=my_input_fn, steps=100)

额外注意事项 #

• 如果你的图数据还包含边信息（比如邻接矩阵），同样可以用上述方法处理，只要在输出签名中标记可变维度即可。
• 使用RaggedTensor时，模型内部需要用tf.ragged相关API处理（比如tf.ragged.mean做全局平均池化得到图的表示）。

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