torch.index_select与torch.gather

旺旺棒棒冰

已于 2022-01-19 10:06:52 修改

阅读量673

点赞数

CC 4.0 BY-SA版权

分类专栏： pytorch 文章标签： pytorch 深度学习 python

于 2022-01-19 10:02:18 首次发布

原文链接：https://pytorch.org/docs/stable/generated/torch.gather.html

pytorch 专栏收录该内容

23 篇文章

订阅专栏

本文详细介绍了PyTorch中两种重要的索引操作：index_select和gather。index_select适用于两维矩阵，能够按行或列选取数据；而gather则能处理多维张量，并按具体索引选取特定位置的数据。

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

torch.index_select

index_select 只能处理两维矩阵，指定行或者列的索引，按行或者按列取出

# indices 只能是一维
x = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [2, 2, 3, 4]])
print(x)
indices = torch.tensor([0, 2])
print(torch.index_select(x, 0, indices))
# 按行取
print(torch.index_select(x, 1, indices))
# 按列取

tensor([[1, 2, 3, 4],
        [5, 6, 7, 8],
        [2, 2, 3, 4]])
# 按行取
tensor([[1, 2, 3, 4],
        [2, 2, 3, 4]])
# 按列取
tensor([[1, 3],
        [5, 7],
        [2, 3]])

torch.gather

gather可处理多维张量，按具体的索引取出对应位置的数（不是整行或者整列取）

x = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [2, 2, 3, 4]])
print(x)
# 3*4



indices = torch.tensor([[1], [2], [3]])
# 3*1
# 意思是取[0,1],[2,2],[3,3]索引位置的数
# 每行取一个，indices指定列标
print(torch.gather(x, 1, indices))
# 按照第二维度取数，indices第一个维度要和x一样


indices = torch.tensor([[1, 2, 1, 0]])
# 1*4
# 意思是取[1,0],[2,1],[1,2],[0,3]索引位置的数
# 每列取一个，indices指定行标
print(torch.gather(x, 0, indices))
# 按照第一维度取数，indices第二个维度要和x一样

tensor([[1, 2, 3, 4],
        [5, 6, 7, 8],
        [2, 2, 3, 4]])
# 每行取一个，指定列标
tensor([[2],
        [7],
        [4]])
# 每列取一个，指定行标
tensor([[5, 2, 7, 4]])