在tensorflow中是否有一个名为“index_select”的pytorch等效函数答案

【问题标题】：Is there an equivalent function of pytorch named "index_select" in tensorflow在tensorflow中是否有一个名为“index_select”的pytorch等效函数
【发布时间】：2019-10-19 14:39:35
【问题描述】：

我试图将 pytorch 代码翻译成 tensorflow。所以我想知道在tensorflow中是否有pytorch的等效函数名为“index_select”

【问题讨论】：

标签： tensorflow pytorch

【解决方案1】：

我还没有找到类似的api可以直接实现，但是我们可以使用tf.slice来实现。

def tf_index_select(input_, dim, indices):
    """
    input_(tensor): input tensor
    dim(int): dimension
    indices(list): selected indices list
    """
    shape = input_.get_shape().as_list()
    if dim == -1:
        dim = len(shape)-1
    shape[dim] = 1

    tmp = []
    for idx in indices:
        begin = [0]*len(shape)
        begin[dim] = idx
        tmp.append(tf.slice(input_, begin, shape))
    res = tf.concat(tmp, axis=dim)

    return res

这是一个显示等价的示例。

import tensorflow as tf
import torch
import numpy as np

a = np.arange(2*3*4).reshape(2,3,4)
dim = 1
indices = [0,2]
# array([[[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11]],

#        [[12, 13, 14, 15],
#         [16, 17, 18, 19],
#         [20, 21, 22, 23]]])

# pytorch
res = torch.tensor(a).index_select(dim, torch.tensor(indices))
# tensor([[[ 0,  1,  2,  3],
#          [ 8,  9, 10, 11]],

#         [[12, 13, 14, 15],
#          [20, 21, 22, 23]]])

# tensorflow
res = tf_index_select(tf.constant(a), dim, indices)
# tensor([[[ 0,  1,  2,  3],
#          [ 8,  9, 10, 11]],

#         [[12, 13, 14, 15],
#          [20, 21, 22, 23]]])

【讨论】：

非常感谢，我发现 tf.gather 也是我需要的，不是吗？
@ShawnChen 是的，但您需要先生成索引。