TensorFlow从列表中收集类似的值

不能得到由不同大小的向量“组”组成的结果。相反,您可以制作一个“group id”张量,根据您的条件将每个向量分类为一个组。使这变得更复杂的部分是你必须用公共元素“融合”组,我认为这只能用一个循环来完成。此代码的作用类似于:

import tensorflow as tf

def make_groups(correspondences):
    # Multiply each row by its index
    m = tf.to_int32(correspondences) * tf.range(tf.shape(correspondences)[0])
    # Pick the largest index for each row
    r = tf.reduce_max(m, axis=1)
    # While loop accounts for transitive correspondences
    # (e.g. if A and B go toghether and B and C go together, then A, B and C go together)
    # The loop makes sure every element gets the largest common group id
    r_prev = -tf.ones_like(r)
    r, _ = tf.while_loop(lambda r, r_prev: tf.reduce_any(tf.not_equal(r, r_prev)),
                         lambda r, r_prev: (tf.gather(r, r), tf.identity(r)),
                         [r, r_prev])
    # Use unique indices to make sequential group ids starting from 0
    return tf.unique(r)[1]

# Test
with tf.Graph().as_default(), tf.Session() as sess:
    arr = tf.constant([[1.5 , 0.2 ],
                       [2.3 , 0.1 ],
                       [1.3 , 0.21],
                       [2.2 , 0.09],
                       [4.4 , 0.8 ],
                       [1.1 , 0.23]])
    a = arr[:, 0]
    b = arr[:, 0]
    cond = (tf.abs(a - a[:, tf.newaxis]) < 0.3) | (tf.abs(b - b[:, tf.newaxis]) < 0.03)
    groups = make_groups(cond)
    print(sess.run(groups))
    # [0 1 0 1 2 0]

因此,在这种情况下,小组将是:

• [1.5, 0.2] , [1.3, 0.21] 和 [1.1, 0.23]
• [2.3, 0.1] 和 [2.2, 0.09]
• [4.4, 0.8]