Question

我有一个形状为（2、4、2）的张量A和一个形状为（4、4）的张量B，所有值均为int。 A中的条目是从0到3。

我想创建一个形状为（2，4，2）的张量C.

for循环代码如下：



for i in range(2):
    for j in range(2):
        for k in range(4):
            C[i][k][j] = B[k][A[i][k][j]]

如何在tensorflow中创建这样的张量C？

谢谢。

Answer 1

您可以使用tf.gather_nd来做到这一点：

import tensorflow as tf

# Input values
A = tf.placeholder(tf.int32, [None, None, None])
B = tf.placeholder(tf.int32, [None, None])
# Make indices for first dimension of B
idx = tf.range(tf.shape(B)[0], dtype=A.dtype)[tf.newaxis, :,  tf.newaxis]
# Tile first dimension indices to match the size of A
idx = tf.tile(idx, (tf.shape(A)[0], 1, tf.shape(A)[2]))
# Stack first dimension indices with A to complete index tensor
idx = tf.stack([idx, A], axis=-1)
# Make result gathering from B
C = tf.gather_nd(B, idx)

这里是一个示例，测试结果是否与您的代码匹配：

import tensorflow as tf
import numpy as np

# Non-TensorFlow implementation for result comparison
A_value = np.random.randint(0, 4, size=(2, 4, 2))
B_value = np.random.randint(100, size=(4, 4))
C_value = np.empty(A_value.shape, dtype=B_value.dtype)
for i in range(A_value.shape[0]):
    for j in range(A_value.shape[2]):
        for k in range(A_value.shape[1]):
            C_value[i][k][j] = B_value[k][A_value[i][k][j]]

# TensorFlow implementation
A = tf.placeholder(tf.int32, [None, None, None])
B = tf.placeholder(tf.int32, [None, None])
idx = tf.range(tf.shape(B)[0], dtype=A.dtype)[tf.newaxis, :,  tf.newaxis]
idx = tf.tile(idx, (tf.shape(A)[0], 1, tf.shape(A)[2]))
idx = tf.stack([idx, A], axis=-1)
C = tf.gather_nd(B, idx)

# Check result
with tf.Session() as sess:
    C_value_tf = sess.run(C, feed_dict={A: A_value, B: B_value})
    print(np.all(np.equal(C_value_tf, C_value)))

输出：

True

如何在张量流中向量化索引操作

1 个答案: