将值求和到由tensorflow中相同形状的张量指定的桶/箱中

Question

我正在寻找一种有效的方法来将浮点变量求和到另一个张量指定的桶中。

我指定有效率是因为实际的输入和输出相当大。它们可以毫无问题地放入内存(数以百万计的元素)，但如果我们将内存或计算复杂性与必要的东西平方，就会遇到问题。

此外，类似c的解决方案只是循环输入并添加到数组/哈希图理论上是有效的，但在TF中导致可怕的执行时间，我正在寻找适当利用多处理的东西。这个要求通常归结为不能对单个数组元素进行非并行循环。

问题示例:

import tensorflow as tf
import numpy as np
buckets = 8
indices = tf.convert_to_tensor([0, 0, 0, 1, 1, 3, 3, 5], tf.int32)
values = tf.convert_to_tensor([.5, .3, .2, .1, 1., 1., 1., .1], tf.float32)
# Inefficient solution that adds a new dimension,
# materializes a dense tensor, and then sums along the added dimension
# memory and computation complexity is O(buckets x indices.size), that's absolutely terrible.
indices_new_axis = tf.range(buckets, dtype=tf.int64)
indices = tf.stack([tf.cast(indices, tf.int64), indices_new_axis], axis=-1)
sparse_repr = tf.SparseTensor(indices, values, dense_shape=[buckets, indices.shape[0]])
dense_repr = tf.sparse.to_dense(sparse_repr)
result = tf.reduce_sum(dense_repr, axis=1)
print(result)
expected_result_dense = [1., 1.1, 0., 2., 0., .1, 0., 0.]
np.testing.assert_array_almost_equal(expected_result_dense, result.numpy())
# the same with sparse representation would also be good:
expected_indices_sparse = [0, 1, 3, 5]
expected_values_sparse = [1., 1.1, 2., .1]

问题的一些技术背景:

我正在尝试对一些具有置信度投票的分析形状进行霍夫变换(没有严格的梯度阈值，但仍然消除了原始图像中的大多数点)。对于没有权重的版本，我只是在索引上使用tf.bincount，这很好地工作，我想知道我是否可以在这里做类似的事情。

我知道，在这个特定的情况下，我可以通过在迭代原始图像中的每个点亮像素时累积结果来完全避免这个问题(不会有重复，所以我可以将投票具体化为buckets大小的密集张量并将它们添加到累加器中)，但这比可能的效率低得多。

Answer 1

一天后，我终于找到了解决办法。这是……嗯…相当简单:

import tensorflow as tf
import numpy as np
buckets = 8
indices = tf.convert_to_tensor([0, 0, 0, 1, 1, 3, 3, 5], tf.int32)
values = tf.convert_to_tensor([.5, .3, .2, .1, 1., 1., 1., .1], tf.float32)
result = tf.math.unsorted_segment_sum(values, indices, buckets)
print(result)
expected_result_dense = [1., 1.1, 0., 2., 0., .1, 0., 0.]
np.testing.assert_array_almost_equal(expected_result_dense, result.numpy())

事实证明，在TF 2.6中已经有一个函数(并且已经在TF中存在了一段时间)可以做到这一点。我找不到它，因为我在tf.bincount周围搜索，这是相似的，但命名完全不同。

在通往上述解决方案的路上，我还发现了另一个使用tf.sparse.reduce_sum()的解决方案(这需要之前将数据更改为SparseTensor)，但上面的速度要快得多。