【pytorch系列】——pytorch版本的batch_gather

最新推荐文章于 2024-01-15 17:17:54 发布

原创最新推荐文章于 2024-01-15 17:17:54 发布 · 1.1k 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能

pytorch 专栏收录该内容

9 篇文章

订阅专栏

本文介绍了一种使用PyTorch实现TensorFlow中tf.batch_gather功能的方法。通过定义一个名为batch_gather的函数，该函数接受两个参数：data和index。此函数将数据转换为NumPy数组，然后遍历索引，从原始数据中收集对应元素，并将结果重新转换为PyTorch张量返回。

部署运行你感兴趣的模型镜像

使用pytorch实现tf.batch_gather：

    def batch_gather(data:torch.Tensor, index:torch.Tensor):
        length = index.shape[0]
        t_index = index.data.numpy()
        t_data = data.data.numpy()
        result = []
        for i in range(length):
            result.append(t_data[i, t_index[i], :])

        return torch.from_numpy(np.array(result))