根据索引的 1D 张量从 3D 张量中选择列（Tensorflow）答案

【问题标题】：Selecting columns from 3D tensor according to a 1D tensor of indices (Tensorflow)根据索引的 1D 张量从 3D 张量中选择列（Tensorflow）
【发布时间】：2019-06-13 21:43:05
【问题描述】：

在给定两个输入的情况下，我正在 tensorflow 中寻找一种方法：

input1，形状为 (batch_size, x, y) 的 3D 张量
input2，形状为(batch_size,) 的一维张量，其值都在[0, y - 1]（含）范围内。

返回形状为(batch_size, x) 的二维张量，使得输出中的ith 元素等于input1 中ith 元素的input2[i]-th 列。 p>

示例：如果input1 = [[[1,2], [3,4]], [[5,6], [7,8]], [[9,10], [11,12]]] （所以input1 的形状是(3, 2, 2)）和

input2 = [0, 1, 1], 那么我想要的输出是[[1,3], [6,8], [10,12]]。

解释：输出中的第0个元素是[1,3]，因为input2中的第0个元素是0；因此，它成为input1 的第 0 个元素中的第 0 列。输出的最后一个元素是[6,8]，因为input2的最后一个元素是1；因此，它成为input1 的最后一个元素中的第一列。

尝试：

我尝试使用 tf.one_hot 来完成此操作，（tf.reduce_sum(input1 * tf.one_hot(input2, y), 2)）但 Tensorflow 在进行乘法运算时变得不高兴，说“ValueError：维度必须相等，但对于 'mul' 是 2 和 3（操作：'Mul ') 输入形状：[3,2,2], [3,2]。"

任何帮助将不胜感激，谢谢！

【问题讨论】：

嗨 Kevin，我是 Google 的一名软件工程师，我想在我们正在撰写的一篇研究论文中包含这个问题的屏幕截图，以展示一些有关 TensorFlow 的代表性问题。如果这对你没问题，你能在许可许可下发布你的问题吗？评论说例如“我根据 Apache License 2.0 许可这个 StackOverflow 问题”对我们来说就足够了。谢谢，大卫
我根据 Apache 许可证 2.0 许可这个 StackOverflow 问题。出于好奇，你能在论文完成后放一个链接吗？祝你好运。
谢谢凯文。该论文现已在https://arxiv.org/abs/2003.09040 上提供。你可以阅读更多关于 TF-Coder 的内容here，也可以自己使用on Colab here。

标签： python tensorflow matrix-multiplication one-hot-encoding

【解决方案1】：

你可以使用tf.map_fn()来实现。

import tensorflow as tf
import numpy as np

input1 = [[[1,2], [3,4]], [[5,6], [7,8]], [[9,10], [11,12]]]
input2 = [0, 1, 1]

tf_input1 = tf.placeholder(shape=(None,2,2),dtype=tf.int32)
tf_input2 = tf.placeholder(shape=(None),dtype=tf.int32)

result = tf.map_fn(lambda x: x[0][:,x[1]], [tf_input1,tf_input2], dtype=tf.int32)

with tf.Session()as sess:
    result = sess.run(result,feed_dict={tf_input1:np.array(input1)
        ,tf_input2:np.array(input2)})
    print(result)

# print
[[ 1  3]
 [ 6  8]
 [10 12]]

编辑

tf.map_fn() 与矢量化操作相比速度较慢。我添加了一个矩阵乘法运算。

# shape= (3,2,1)
result = tf.cast(tf.expand_dims(tf.one_hot(input2, 2),-1),tf.int32)
# shape= (3,2)
result = tf.squeeze(tf.matmul(tf_input1, result))

【讨论】：

效率如何？它和矩阵乘法一样高效吗？
@Kevin 我将其添加到答案中。