pytorch中index_select()的用法

最新推荐文章于 2025-07-21 00:28:47 发布

原创最新推荐文章于 2025-07-21 00:28:47 发布 · 5.2w 阅读

169

171 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习入门

pytorch 专栏收录该内容

1 篇文章

订阅专栏

本文详细解析了PyTorch中index_select()函数的使用方法，通过具体实例展示了如何通过该函数按行或列选取张量中的元素。文章首先创建了一个12元素的张量，并将其重塑为3x4的矩阵，然后使用index_select()函数选取特定行和列，清晰地解释了参数的意义。

pytorch中index_select()的用法

刚开始学习pytorch，遇到了index_select(),一开始不太明白几个参数的意思，后来查了一下资料，算是明白了一点。


a = torch.linspace(1, 12, steps=12).view(3, 4)
print(a)
b = torch.index_select(a, 0, torch.tensor([0, 2]))
print(b)
print(a.index_select(0, torch.tensor([0, 2])))
c = torch.index_select(a, 1, torch.tensor([1, 3]))
print(c)

先定义了一个tensor，这里用到了linspace和view方法。
第一个参数是索引的对象，第二个参数0表示按行索引，1表示按列进行索引，第三个参数是一个tensor，就是索引的序号，比如b里面tensor[0， 2]表示第0行和第2行，c里面tensor[1, 3]表示第1列和第3列。

输出结果如下：

tensor([[ 1.,  2.,  3.,  4.],
        [ 5.,  6.,  7.,  8.],
        [ 9., 10., 11., 12.]])
tensor([[ 1.,  2.,  3.,  4.],
        [ 9., 10., 11., 12.]])
tensor([[ 1.,  2.,  3.,  4.],
        [ 9., 10., 11., 12.]])
tensor([[ 2.,  4.],
        [ 6.,  8.],
        [10., 12.]])