为 2D 张量创建索引

Question

我正在尝试做这样的事情：

假设输入张量是(2, 3)张量，其值如下：

[[1 2]
 [3 4]
 [5 6]]

我正在通过以下方式将 2D 张量展平为 1D：

input = tf.reshape(input, [-1])

所以输入现在变成[1 2 3 4 5 6]。

但我还创建了一个索引先前张量的一维张量，因此所需的输出是[0 0 1 1 2 2]。我应该如何在 TF 中创建这个 Tensor？

一般来说，如果输入张量的形状为(X, Y)。我想创建一个看起来像这样的一维张量：

[0 0 0 ... 0 1 ....1 2 ... 2 ... X-1 ... X-1]

每个值重复 Y - 1 次。

Answer 1

这是一种方法；这样做是首先创建一个单行索引，然后使用 tf.tile 按原始张量具有的列数重复每个索引；重塑二维索引可以满足您的需求。

t = tf.constant([[1,2],[3,4],[5,6]])

X, Y = t.shape
idx = tf.range(X.value)
idx_2d = tf.reshape(idx, [-1,1])
idx_2d_full = tf.tile(idx_2d, [1, Y.value])
idx_flat = tf.reshape(idx_2d_full, [-1])

with tf.Session() as sess:
    print(sess.run(idx_flat))

[0 0 1 1 2 2]