如何使用张量索引张量流中的另一个张量 [英] How to use a tensor for indexing another tensor in tensorflow

问题描述

我有一个 data 维度张量 [BXNX 3]，我有一个 indices 维度张量 [BXM].我希望使用 indices 张量从 data 张量中提取 [B X M X 3] 张量.

I have a data tensor of dimensios [B X N X 3], and I have an indices tensor of dimensions [B X M]. I wish to extract a [B X M X 3] tensor from the data tensor using the indices tensor.

我有这个有效的代码:

new_data= []    
for i in range(B):
        new_data.append(tf.gather(data[i], indices[i]))
new_data= tf.stack(new_data) 

但是，我确信这不是正确的方法.有人知道更好的方法吗?(我想我应该以某种方式使用 tf.gather_nd() 但我不知道如何使用)

However, I am sure it is not the right way to do this. Does anyone know a better way? (I guess I should use tf.gather_nd() somehow but I couldn't figure out how)

我在这里看到了几个类似问题的答案.但是我找不到解决我的问题的方法.

I have seen several answers to similar questions here. However I could not find the solution to my problem.

推荐答案

你可以使用 tf.gather_nd() 和这样的代码:

You can use tf.gather_nd() with code like this:

import tensorflow as tf

# B = 3
# N = 4
# M = 2
# [B x N x 3]
data = tf.constant([
    [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]],
    [[100, 101, 102], [103, 104, 105], [106, 107, 108], [109, 110, 111]],
    [[200, 201, 202], [203, 204, 205], [206, 207, 208], [209, 210, 211]],
    ])

# [B x M]
indices = tf.constant([
    [0, 2],
    [1, 3],
    [3, 2],
    ])

indices_shape = tf.shape(indices)

indices_help = tf.tile(tf.reshape(tf.range(indices_shape[0]), [indices_shape[0], 1]) ,[1, indices_shape[1]]);
indices_ext = tf.concat([tf.expand_dims(indices_help, 2), tf.expand_dims(indices, 2)], axis = 2)
new_data = tf.gather_nd(data, indices_ext)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print('data')
    print(sess.run(data))
    print('\nindices')
    print(sess.run(indices))
    print('\nnew_data')
    print(sess.run(new_data))

new_data 将是:

[[[  0   1   2]
  [  6   7   8]]

 [[103 104 105]
  [109 110 111]]

 [[209 210 211]
  [206 207 208]]]

这篇关于如何使用张量索引张量流中的另一个张量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！