如何根据张量对 Tensorflow 中的值进行索引？

Question

我在 tensorflow v1 下面给出了三个张量：

mod_labels = [0,1,1,0,0]
feats = [[1,2,1], [3,2,6], [1,1,1], [9,8,4], [5,4,8]]
labels = [1,53,12,89,54]

我想根据mod_labels中的值创建四个新张量：

# Make new tensors for mod_labels=0
mod0_feats = [[1,2,1], [9,8,4], [5,4,8]]
mod0_labels = [1,89,54]

# Similarly make tensors for mod_labels=1
mod1_feats = [[3,2,6], [1,1,1]]
mod1_labels = [53,12]

我尝试使用 for 循环遍历 mod_labels 但 tensorflow 不允许遍历占位符。

Answer 1

好吧，假设这些是张量：

mod_labels = tf.convert_to_tensor([0,1,1,0,0])
feats = tf.convert_to_tensor([[1,2,1], [3,2,6], [1,1,1], [9,8,4], [5,4,8]])
labels = tf.convert_to_tensor([1,53,12,89,54])
zeros = tf.where(mod_labels == 0)
ones = tf.where(mod_labels == 1)
mod0_feats = tf.gather_nd(feats, zeros)
mod0_labels = tf.gather_nd(labels, zeros)
mod1_feats = tf.gather_nd(feats, ones)
mod1_labels = tf.gather_nd(labels, ones)