如何在張量流中獲得索引張量？

Question

例如，我有一個形狀為 [2, 2] 的張量。 我想得到一個索引的 3D 張量：[[[0, 0], [0, 1]], [[1, 0], [1, 1]]]。

我正在使用以下方法：

my_shape = [2, 2]
my_range = my_shape[0]*my_shape[1]
m = tf.range(0, my_range)
def get_inds(t, last_dim):
  return tf.convert_to_tensor([t // last_dim, t % last_dim])

inds = tf.map_fn(fn=lambda t: get_inds(t, my_shape[-1]), elems=m)
sh = tf.concat([my_shape, [2]], -1)
inds = tf.reshape(inds, sh)

有沒有更好的方法？

Answer 1

我想在 numpy 中這樣做： meshgrid可能是要走的路（結合stack以獲得 1 個數組）。

import tensorflow as tf
shape = (2,2)
x = tf.range(shape[0])
y = tf.range(shape[1])
xx,yy = tf.meshgrid(x,y)
indices = tf.stack([yy,xx],axis=2)

>>> indices
<tf.Tensor: shape=(2, 2, 2), dtype=int32, numpy=
array([[[0, 0],
        [0, 1]],

       [[1, 0],
        [1, 1]]], dtype=int32)>