Pytorch：tensor.gather

最新推荐文章于 2024-05-29 01:07:50 发布

X-ocean

最新推荐文章于 2024-05-29 01:07:50 发布

阅读量127

点赞数

文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/xieocean/article/details/131939322

版权

该文通过一个例子解释了PyTorch中gather函数的作用，特别是dim参数的含义。输入是一个3x3的张量，index是一个二维索引tensor，通过在dim=1的维度上应用索引，将张量的对应位置值取出，形成了一个新的输出张量。最终结果是将原张量的第三列（索引2）移到了第一列，第二列（索引1）移到了第二列，第一列（索引0）移到了第三列，从而展示了gather函数如何根据提供的索引获取张量的值。

摘要由CSDN通过智能技术生成

这个函数不好理解，但是在知乎上看到一个言简意赅的解释，特地记录。

例子：

input= torch.arange(3, 12).view(3, 3)
# tensor([[ 3,  4,  5],
#        [ 6,  7,  8],
#        [ 9, 10, 11]])
index = torch.tensor([[2, 1, 0]])
output= input.gather(dim=1, index)

解释：

首先明确一点，index和output的shape是一样的，即index.shape = output.shape。然后可以确定output = [ [ ?, ?, ?] ]
接着把output写成索引形式，即output = [ [ (0,1), (0,2), (0,3) ] ]
然后，由于gather中的dim参数是1，我们就把索引形式的output在1这个维度上替换成index的元素，即output 从 [ [ (0,1), (0,2), (0,3) ] ] 变成了 [ [ (0,2), (0,1), (0,0) ] ]
然后按照索引取元素，得到output =[ [ 5, 4, 3] ]