如何在tf.data.Dataset中使用稀疏张量列表？

Question

我正在尝试构建一个以稀疏张量列表作为输入的模型。(列表长度等于批量大小)

我使用稀疏张量的原因是我必须将邻接矩阵传递给我的GNN模型，而它是非常稀疏的。(~99%)

我很熟悉pytorch的使用，它很容易将稀疏张量输入到网络中。

然而，我发现我必须使用tf.data.Dataset或keras.utils.Sequence在tensorflow中创建数据集。

但是当我使用稀疏张量列表作为输入时，这些方法会抛出错误。

例如，下面的代码使TypeError

import tensorflow as tf
tf.data.Dataset.from_tensor_slices(sparse_lists)

TypeError: Neither a SparseTensor nor SparseTensorValue: 
[<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2e25b5c0>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c22ada0>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c22a400>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c1ed240>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c1ed390>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c1ed470>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c1ed5c0>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c1ed710>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c1ed828>, 
<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fbf2c1ed940>].

我知道，如果我将列表中的所有稀疏张量连接成一个巨大的张量，它就会起作用。然而，这不是我的选择，因为稍后我必须对稀疏张量使用索引。(如果我将2D稀疏张量合并为3D稀疏张量，我不能像下面这样使用索引)

Some3DSparseTensor[:10]

此外，这将需要更多的时间，因为我必须切片3D张量，以便与其他密集网络进行矩阵乘法。

此外，我知道如果我为每一批做稀疏张量是可以的，但是每一批会花费太多的时间。

因此，我想让tf.data.Dataset能够从稀疏张量列表中生成批处理，因为索引和时间问题。

有人能帮我吗？:)

长话短说，

我有:稀疏张量列表(例如1000000长度列表)

我需要做的是:稀疏张量的批量列表(例如1024长度列表，而不是稀疏连接)

Answer 1

如果SparseTensors具有相同的dense_shape，您可以创建一个惟一的SparseTensor而不是列表，并将其传递给from_tensor_slices。

例如，下面的代码从一个大的SparseTensor s中产生单独的SparseTensors，沿着第一个维度将它们分开

s = tf.sparse.SparseTensor(
    indices=tf.constant([[0, 0, 0], [1, 0, 0], [1, 0, 1], [2, 1, 1]], dtype=tf.int64), 
    values=tf.range(4, dtype=tf.float32), 
    dense_shape=(3, 2, 2))
d = tf.data.Dataset.from_tensor_slices(s)
for t in d:
    print(t)
>>> SparseTensor(indices=tf.Tensor([[0 0]], shape=(1, 2), dtype=int64), values=tf.Tensor([0.], shape=(1,), dtype=float32), dense_shape=tf.Tensor([2 2], shape=(2,), dtype=int64))
SparseTensor(indices=tf.Tensor(
[[0 0]
 [0 1]], shape=(2, 2), dtype=int64), values=tf.Tensor([1. 2.], shape=(2,), dtype=float32), dense_shape=tf.Tensor([2 2], shape=(2,), dtype=int64))
SparseTensor(indices=tf.Tensor([[1 1]], shape=(1, 2), dtype=int64), values=tf.Tensor([3.], shape=(1,), dtype=float32), dense_shape=tf.Tensor([2 2], shape=(2,), dtype=int64))

要以这种方式使用from_tensor_slices，您需要一个函数将列表sparse_lists转换为较大的SparseTensor s (如下所述)。

简单地说，你可以这样做

import tensorflow as tf

def sparse_list_to_sparse_tensor(sparse_lists):
    n = len(sparse_lists)
    shape = sparse_lists[0].dense_shape
    out_shape = (n, *shape)
    out_values = tf.concat([s.values for s in sparse_lists], axis=0)
    out_indices = []
    for i, s in enumerate(sparse_lists):
        element_idx = tf.cast(tf.fill((s.indices.shape[0], 1), i), dtype=tf.int64)
        out_indices.append(tf.concat([element_idx, s.indices], axis=1))
    out_indices = tf.concat(out_indices, axis=0)
    return tf.sparse.SparseTensor(out_indices, out_values, out_shape)


tf.data.Dataset.from_tensor_slices(sparse_list_to_sparse_tensor(sparse_lists))

Answer 2

另一种解决方案是在每个稀疏张量上使用from_tensor_slices (在添加虚拟批次维度之后)，以创建具有单个元素的多个数据集，这些元素可以连接到单个数据集中。

dataset = None
for sparse_tensor in sparse_list:
    batched_sparse_tensor = tf.sparse.expand_dims(sparse_tensor, axis=0)
    element_dataset = tf.data.Dataset.from_tensor_slices(batched_sparse_tensor)
    if dataset is None:
        dataset = element_dataset
    else:
        dataset = dataset.concatenate(element_dataset)

请注意，使用此解决方案，稀疏张量可以具有不同的dense_shape。