python - tensorflow collect similar values from a list

Question

I have a tensor as follows:

arr = [[1.5,0.2],[2.3,0.1],[1.3,0.21],[2.2,0.09],[4.4,0.8]]

I would like to collect small arrays whose difference of first elements are within 0.3 and second elements are within 0.03. For example [1.5,0.2] and [1.3,0.21] should belong to a same category. The difference of their first elements is 0.2<0.3 and second 0.01<0.03.

I want a tensor looks like this

arr = {[[1.5,0.2],[1.3,0.21]],[[2.3,0.1],[2.2,0.09]]}

How to do this in tensorflow? Eager mode is ok.

I found a way which is a bit ugly and slow:

samples = np.array([[1.5,0.2],[2.3,0.1],[1.3,0.2],[2.2,0.09],[4.4,0.8],[2.3,0.11]],dtype=np.float32)
ini_samples = samples
samples = tf.split(samples,2,1)

a = samples[0]
b = samples[1]

find_match1 = tf.reduce_sum(tf.abs(tf.expand_dims(a,0) - tf.expand_dims(a,1)),2)
a = tf.logical_and(tf.greater(find_match1, tf.zeros_like(find_match1)),tf.less(find_match1, 0.3*tf.ones_like(find_match1)))

find_match2 = tf.reduce_sum(tf.abs(tf.expand_dims(b,0) - tf.expand_dims(b,1)),2)
b = tf.logical_and(tf.greater(find_match2, tf.zeros_like(find_match2)),tf.less(find_match2, 0.03*tf.ones_like(find_match2)))
x,y = tf.unique(tf.reshape(tf.where(tf.logical_or(a,b)),[1,-1])[0])

r = tf.gather(ini_samples, x)

Does tensorflow have more elegant functions?

标签： pythontensorflow

Answer 1

You cannot get a result composed of "groups" of vectors with different sizes. Instead, you can make a "group id" tensor that classifies each vector into a group according to your criteria. The part that makes this a bit more complicated is that you have to "fuse" groups with common elements, which I think can only be done with a loop. This code does something like that:

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]

So in this case, the groups would be:

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