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

最新推荐文章于 2024-09-19 21:01:33 发布

黑白象

最新推荐文章于 2024-09-19 21:01:33 发布

阅读量1.1k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： pytorch 文章标签：深度学习人工智能

本文链接：https://blog.youkuaiyun.com/dendi_hust/article/details/103957914

pytorch 专栏收录该内容

9 篇文章

订阅专栏

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

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

使用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))