gather torch_gather函数

Python

Python开发

Python语言

gather函数

gather(input, dim, index):根据  index,在  dim  维度上选取数据,输出的  size  与  index  一致

# input (Tensor) – 源张量

# dim (int) – 索引的轴

# index (LongTensor) – 聚合元素的下标(index需要是torch.longTensor类型)

# out (Tensor, optional) – 目标张量

for 3D tensor:

out[i][j][k] = tensor[index[i][j][k]][j][k]   # dim=0

out[i][j][k] = tensor[i][index[i][j][k]][k]   # dim=1

out[i][j][k] = tensor[i][j][index[i][j][k]]   # dim=2

for 2D tensor:

out[i][j] = input[index[i][j]][j]  # dim = 0

out[i][j] = input[i][index[i][j]]  # dim = 1

import torch as t # 导入torch模块

c = t.arange(0, 60).view(3, 4, 5) # 定义tensor

print(c)

index = torch.LongTensor([[[0,1,2,0,2],

[0,0,0,0,0],

[1,1,1,1,1]],

[[1,2,2,2,2],

[0,0,0,0,0],

[2,2,2,2,2]]])

b = t.gather(c, 0, index)

print(b)

输出:

tensor([[[ 0, 1, 2, 3, 4],

[ 5, 6, 7, 8, 9],

[10, 11, 12, 13, 14],

[15, 16, 17, 18, 19]],

[[20, 21, 22, 23, 24],

[25, 26, 27, 28, 29],

[30, 31, 32, 33, 34],

[35, 36, 37, 38, 39]],

[[40, 41, 42, 43, 44],

[45, 46, 47, 48, 49],

[50, 51, 52, 53, 54],

[55, 56, 57, 58, 59]]])

报错:

Traceback (most recent call last):

File "E:/Release02/my_torch.py", line 14, in

b = t.gather(c, 0, index)

RuntimeError: Size does not match at dimension 1 get 4 vs 3

(第1维尺寸不匹配)

将index调整为:

index = t.LongTensor([[[0, 1, 2, 0, 2], [0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]],

[[1, 2, 2, 2, 2], [0, 0, 0, 0, 0], [2, 2, 2, 2, 2], [1, 1, 1, 1, 1]],

[[1, 2, 2, 2, 2], [0, 0, 0, 0, 0], [2, 2, 2, 2, 2], [1, 1, 1, 1, 1]]])

则上文输出为:

tensor([[[ 0, 21, 42, 3, 44],

[ 5, 6, 7, 8, 9],

[30, 31, 32, 33, 34],

[35, 36, 37, 38, 39]],

[[20, 41, 42, 43, 44],

[ 5, 6, 7, 8, 9],

[50, 51, 52, 53, 54],

[35, 36, 37, 38, 39]],

[[20, 41, 42, 43, 44],

[ 5, 6, 7, 8, 9],

[50, 51, 52, 53, 54],

[35, 36, 37, 38, 39]]])

对于2D tensor 则无“index与tensor 的size一致”之要求,

这个要求在官方文档和其他博文、日志中均无提到

(可能是个坑吧丨可能是个坑吧丨可能是个坑吧)

eg:

代码(此部分来自https://www.yzlfxy.com/jiaocheng/python/337618.html):

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]

官方文档:

torch.gather(input, dim, index, out=None) → Tensor

Gathers values along an axis specified by dim.

For a 3-D tensor the output is specified by:

out[i][j][k] = input[index[i][j][k]][j][k] # dim=0

out[i][j][k] = input[i][index[i][j][k]][k] # dim=1

out[i][j][k] = input[i][j][index[i][j][k]] # dim=2

Parameters:

input (Tensor)-The source tensor

dim (int)-The axis along which to index

index (LongTensor)-The indices of elements to gather

out (Tensor, optional)-Destination tensor

Example:

>>> t = torch.Tensor([[1,2],[3,4]])

>>> torch.gather(t, 1, torch.LongTensor([[0,0],[1,0]]))

1 1

4 3

[torch.FloatTensor of size 2x2]

以上,学习中遇到的问题,记录方便回顾,亦示他人以之勉坑

内容来源于网络,如有侵权请联系客服删除

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值