如何在tf.data.Dataset中使用稀疏张量列表构建数据集？

解决方案 #

根因说明 #

tf.data.Dataset.from_tensor_slices 不支持直接传入多个SparseTensor组成的列表，仅支持输入单个SparseTensor或者可堆叠的稠密张量结构，因此会触发你遇到的类型错误。

实现方案 #

使用tf.data.Dataset.from_generator构造数据集，不需要提前拼接所有稀疏张量，也不需要逐批次重构SparseTensor，几乎没有额外性能开销，完全符合你的需求。

完整代码示例

import tensorflow as tf

# --------------- 以下为测试数据构造，替换为你自己的sparse_lists即可 ---------------
sparse_lists = []
# 模拟10个2D稀疏邻接矩阵，形状为100x100，按需修改为你实际的张量形状和数据类型
for i in range(10):
    indices = tf.random.uniform((10, 2), minval=0, maxval=100, dtype=tf.int64)
    values = tf.random.normal((10,))
    sparse_lists.append(tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=(100, 100)))
# --------------------------------------------------------------------------------

# 1. 定义生成器，直接遍历稀疏张量列表返回单个元素
def sparse_generator():
    for st in sparse_lists:
        yield st

# 2. 构造数据集，指定输出签名，shape和dtype替换为你实际的稀疏张量参数
dataset = tf.data.Dataset.from_generator(
    sparse_generator,
    output_signature=tf.SparseTensorSpec(shape=(100, 100), dtype=tf.float32)
)

batch_size = 1024 # 按需修改批次大小

# ---------------- 两种批次输出方式二选一即可 ----------------
# 方式1：直接输出长度为batch_size的独立稀疏张量列表，完全无拼接
for window in dataset.window(batch_size, drop_remainder=True):
    batch = list(window.as_numpy_iterator())
    # 此处batch即为你需要的批次数据，每个元素是原始的2D SparseTensor
    # 后续可直接传入GNN模型使用

# 方式2：输出为堆叠的3D SparseTensor，支持原生索引操作
# dataset = dataset.batch(batch_size, drop_remainder=True)
# for batch in dataset:
    # 支持batch[:10]、batch[i]等索引操作，返回的是对应位置的2D SparseTensor
    # 索引效率和原生列表一致，无额外性能开销

方案优势 #

• 无额外预处理开销：直接复用原始的稀疏张量对象，不需要拼接或者重构稀疏结构
• 兼容tf.data全链路操作：可叠加shuffle、prefetch、cache等流水线优化，不影响训练效率
• 适配不同使用习惯：支持输出独立稀疏张量列表或者可索引的堆叠稀疏张量两种格式，按需选择即可

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