Как эффективно группировать тензоры в тензорном потоке

Question

Я ищу эффективный способ сгруппировать индексы набора тензорных данных в тензор потока 2, python.

Ниже приведен пример действия для выполнения

Input
------------------
id | feature_frame

 0 | [0, 1, 2, 3]
 1 | [0, 1, 2, 3]
 2 | [3, 1, 1, 2]
 3 | [7, 7, 1, 2]
 4 | [0, 1, 2, 3]
 5 | [3, 1, 1, 2]

Output
------------------
0 | [0,1,4]
1 | [2,5]
2 | [3]

Как видите, индексы, соответствующие равным тензорам, сгруппированы вместе. Порядок, в котором они сгруппированы и где расположена каждая группа, не имеет значения.

Answer 1

С целочисленными значениями, если число столбцов и значений не слишком велико, вы можете вычислить идеальное га sh для каждого характерного кадра, то есть одно целое число, которое однозначно представляет каждую комбинацию значений. Тогда это довольно легко решить.

import tensorflow as tf

# Assumes values in f are non-negative, you could shift by minimum if not
f = tf.constant([[0, 1, 2, 3],
                 [0, 1, 2, 3],
                 [3, 1, 1, 2],
                 [7, 7, 1, 2],
                 [0, 1, 2, 3],
                 [3, 1, 1, 2]], dtype=tf.int32)

# Compute hash
m = tf.reduce_max(f, axis=0)
b = tf.math.cumprod(m, exclusive=True)
h = tf.reduce_sum(f * b, axis=1)
# Find unique positions and counts
_, idx, counts = tf.unique_with_counts(h)
# Make groups
groups = tf.argsort(idx)
print(groups.numpy())
# [0 1 4 2 5 3]
# The counts array gives the size of each consecutive groups
print(counts.numpy())
# [3 2 1]
# Or you can get the array of group indices
group_idx = tf.gather(idx, groups)
print(group_idx.numpy())
# [0 0 0 1 1 2]
# Or even a ragged tensor
groups_rag = tf.RaggedTensor.from_row_lengths(groups, counts)
print(groups_rag.to_list())
# [[0, 1, 4], [2, 5], [3]]

Если вы не можете использовать трюк хеширования, вы всегда можете сравнить каждую строку друг с другом, что дороже, но всегда должно работать:

import tensorflow as tf

f = tf.constant([[0, 1, 2, 3],
                 [0, 1, 2, 3],
                 [3, 1, 1, 2],
                 [7, 7, 1, 2],
                 [0, 1, 2, 3],
                 [3, 1, 1, 2]], dtype=tf.int32)

# All to all comparison
cmp = tf.equal(tf.expand_dims(f, axis=1), f)
# Find matching rows
eq = tf.reduce_all(cmp, axis=-1)
# Unfortunately tf.math.argmax is indeterministic on ties
# so we multiply by an index array to find first match for each row
r = tf.range(tf.shape(f)[0], 0, -1)
match_idx = r * tf.dtypes.cast(eq, r.dtype)
group_id = tf.math.argmax(match_idx, axis=1)
# Continue as before
_, idx, counts = tf.unique_with_counts(group_id)
# ...

РЕДАКТИРОВАТЬ: еще одна альтернатива с использованием строк, как это было предложено:

import tensorflow as tf

f = tf.constant([[0, 1, 2, 3],
                 [0, 1, 2, 3],
                 [3, 1, 1, 2],
                 [7, 7, 1, 2],
                 [0, 1, 2, 3],
                 [3, 1, 1, 2]], dtype=tf.int32)

# Convert to string
fs = tf.strings.as_string(f)
# Join with some non-numerical character
h = tf.strings.reduce_join(fs, axis=1, separator=' ')
# Continue as before
_, idx, counts = tf.unique_with_counts(h)
# ...

РЕДАКТИРОВАТЬ: Ради полноты, вот еще одна версия, использующая вложенный l oop, вероятно, не очень быстрый:

import tensorflow as tf

f = tf.constant([[0, 1, 2, 3],
                 [0, 1, 2, 3],
                 [3, 1, 1, 2],
                 [7, 7, 1, 2],
                 [0, 1, 2, 3],
                 [3, 1, 1, 2]], dtype=tf.int32)

n = tf.shape(f)[0]
ta = tf.TensorArray(tf.int32, size=n, element_shape=())
_, ta = tf.while_loop(lambda i, ta: i < n,
                      lambda i, ta: (i + 1, ta.write(i, tf.while_loop(
                          lambda j: ~tf.reduce_all(tf.equal(f[i], f[j])),
                          lambda j: [j + 1],
                          [0])[0])),
                      [0, ta])
h = ta.stack(0)
# Continue as before
_, idx, counts = tf.unique_with_counts(h)
# ...