tf.math.segment_sum 的 numpy 等价物

社会演员多 2年前 tensorflow 344

原文标题 ：numpy equivalent of tf.math.segment_sum

什么是等效的 tf.math.segment_sumin numpy？所以基本上我喜欢将 tf 中完全相同的代码重写为 np，其中我使用段 sum 使用 segment_ids 数组将某些元素组合在一起并对这些段求和。什么是等效代码在麻木？我有一个数组和segment_ids 数组，我喜欢在numpy 中执行segment_sum。

原文链接：https://stackoverflow.com//questions/71681846/numpy-equivalent-of-tf-math-segment-sum

我来回复

Warren Weckesser 评论

您可以使用方法numpy.add.at创建非常接近tf.math.segment_sum的东西，这是addufunc的at方法：

def segment_sum(data, segment_ids):
    data = np.asarray(data)
    s = np.zeros((np.max(segment_ids)+1,) + data.shape[1:], dtype=data.dtype)
    np.add.at(s, segment_ids, data)
    return s

例如，

In [53]: c = np.array([[1, 2, 3, 4], [4, 3, 2, 1], [5, 6, 7, 8]])

In [54]: ids = [0, 0, 1]

In [55]: segment_sum(c, ids)
Out[55]: 
array([[5, 5, 5, 5],
       [5, 6, 7, 8]])

In [56]: x = [10, 20, 20, 30, 10, 0, 1, 2]

In [57]: xids = [1, 1, 0, 0, 2, 2, 2, 3]

In [58]: segment_sum(x, xids)
Out[58]: array([50, 30, 11,  2])

In [59]: w = np.arange(72).reshape(6, 2, 6) % 5

In [60]: w
Out[60]: 
array([[[0, 1, 2, 3, 4, 0],
        [1, 2, 3, 4, 0, 1]],

       [[2, 3, 4, 0, 1, 2],
        [3, 4, 0, 1, 2, 3]],

       [[4, 0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4, 0]],

       [[1, 2, 3, 4, 0, 1],
        [2, 3, 4, 0, 1, 2]],

       [[3, 4, 0, 1, 2, 3],
        [4, 0, 1, 2, 3, 4]],

       [[0, 1, 2, 3, 4, 0],
        [1, 2, 3, 4, 0, 1]]])

In [61]: wids = [0, 0, 1, 2, 2, 2]

In [62]: segment_sum(w, wids)
Out[62]: 
array([[[2, 4, 6, 3, 5, 2],
        [4, 6, 3, 5, 2, 4]],

       [[4, 0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4, 0]],

       [[4, 7, 5, 8, 6, 4],
        [7, 5, 8, 6, 4, 7]]])

2年前 0条评论

tf.math.segment_sum 的 numpy 等价物

回复

相关问题