Pytorch中torch.gather函数_torch.arrange()函数用法 - CSDN文库

5星 · 超过95%的资源 16 浏览量更新于2023-10-24 评论收藏 35KB PDF 举报

身份认证购VIP最低享 7 折!

领优惠券(最高得80元）

资源详情

资源评论

资源推荐

Pytorch中中torch.gather函数函数

在学习 CS231n中的NetworkVisualization-PyTorch任务，讲解了使用torch.gather函数，gather函数是用来根据你输入的位置

索引 index，来对张量位置的数据进行合并，然后再输出。

其中 gather有两种使用方式，一种为 torch.gather

另一种为对象.gather。

首先介绍对象.gather

import torch

torch.manual_seed(2) #为CPU设置种子用于生成随机数，以使得结果是确定的

def gather_example():

N, C = 4, 5

s = torch.randn(N, C)

y = torch.LongTensor([1, 2, 1, 3]) # 必须要为 LongTensor 不然会报错

print(s)

print(y)

print(s.gather(1,y.view(-1, 1)).squeeze())

gather_example()

'''

输出:

tensor([[-1.0408, 0.9166, -1.3042, -1.1097, 0.0299],

[-0.0498, 1.0651, 0.8860, -0.8110, 0.6737],

[-1.1233, -0.0919, 0.1405, 1.1191, 0.3152],

[ 1.7528, -0.7396, -1.2425, -0.1752, 0.6990]])

tensor([1, 2, 1, 3])

tensor([ 0.9166, 0.8860, -0.0919, -0.1752])

'''

对于上图的代码,首先通过 torch.randn 随机输出化出结果为

tensor([[-1.0408, 0.9166, -1.3042, -1.1097, 0.0299],

[-0.0498, 1.0651, 0.8860, -0.8110, 0.6737],

[-1.1233, -0.0919, 0.1405, 1.1191, 0.3152],

[ 1.7528, -0.7396, -1.2425, -0.1752, 0.6990]])

然后我们根据索引 tensor([1, 2, 1, 3]) 对每一行进行索引，在第0行索引到位置=1的元素，即 0.9166，在第二行索引到位置=2

的元素即 0.8860 以此类推，即为最后的结果。

另一种为 torch.gather

b = torch.Tensor([[1,2,3],[4,5,6]])

print (b)

index_1 = torch.LongTensor([[0,1],[2,0]])

index_2 = torch.LongTensor([[0,1,1],[0,0,0]])

print (torch.gather(b, dim=1, index=index_1)) # 按行来进行索引

print (torch.gather(b, dim=0, index=index_2)) # 按列来进行索引

'''

输出为：

1 2 3

4 5 6

[torch.FloatTensor of size 2x3]

1 2

6 4

[torch.FloatTensor of size 2x2]

1 5 6

1 2 3

[torch.FloatTensor of size 2x3]

'''

当dim=0时候，那么就是按照行进行索引，输出第一行位置为0,1的元素，即1,2 。第二行位置为2,0的元素，即 6,4。

本内容试读结束，登录后可阅读更多

下载后可阅读完整内容，剩余1页未读，立即下载

评论10

番皂泡

2023-06-16

挺不错的一篇文章，让我学到了很多关于torch.gather函数的知识。

weixin_38612811

粉丝: 5
资源: 933

会员权益专享

图片转文字

全年可省5，000元立即开通

最新资源

资源上传下载、课程学习等过程中有任何疑问或建议，欢迎提出宝贵意见哦~我们会及时处理！点击此处反馈