高级操作---where 和 gather

最新推荐文章于 2024-07-12 17:01:23 发布

0iq2333

最新推荐文章于 2024-07-12 17:01:23 发布

阅读量1.1k

点赞数

CC 4.0 BY-SA版权

分类专栏：从零开始学深度学习文章标签： pytorch 深度学习 python

本文链接：https://blog.youkuaiyun.com/acm_durante/article/details/122729957

从零开始学深度学习专栏收录该内容

18 篇文章

订阅专栏

本文详细介绍了PyTorch中的两个高级函数：where和gather。where函数可以根据条件选择不同的tensor元素，而gather函数则用于从指定维度上收集tensor的元素。文章通过具体的代码示例展示了这两个函数的应用。

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

一、where函数

1、简介

return a tensor of elements selected from either x or y, depending on condition
返回一个tensor，根据生成规则选择是从x取出还是y取出。
虽然where函数操作可以用for循环代替，但是where是通过GPU加速的，在深度学习中使用where速度会更快。

2.函数

torch.where(cond,x,y)
cond:生成规则自定义
x,y:两个tensor

3.代码实例

cond = torch.rand(2,2) # 倾向于a的概率
print(cond)
# tensor([[0.7894, 0.3648],
#        [0.4843, 0.3184]])

a = torch.full([2,2],0)

b = torch.full([2,2],1)
# 规则是如果cond>0.5就选择a，否则选择b
ans = torch.where(cond>0.5,a,b)
print(ans)
# tensor([[0., 1.],
#        [1., 1.]])

二、gather函数

1、简介

官方文档：
1. torch.gather(input,dim,index,out=None) -> Tensor
2. Gather values along an axis specified by dim
3. For a 3D tensor the output is specified by:
  1. out[i][j][k] = input[ index[i][j][k] ][j][k], dim = 0
  2. out[i][j][k] = input[i][ index[i][j][k] ][k], dim = 1
  3. out[i][j][k] = input[i][j][ index[i][j][k] ] , dim = 2
从原tensor中获取指定dim和指定index的数据

2、函数

torch.gather(input,dim,index,out=None)
input:保持与index.shape一致，从input中获取数值
dim:从input中的哪一维度获取
index:保持与input.shape一致，可以理解为tensor下标，根据index生成最终的tensor

3、代码实例

prob = torch.randn(4,10) # 随机生成shape为[4,10]的tensor
idx = prob.topk(dim=1,k=3) # 生成前3大的tensor
idx = idx[1]
print(idx)
# tensor([[8, 7, 9],
#        [1, 3, 2],
#        [3, 1, 6],
#        [7, 8, 9]])

label = torch.arange(10) + 100 # 偏移
print(label)
# tensor([100, 101, 102, 103, 104, 105, 106, 107, 108, 109])

ans = torch.gather(label.expand(4,10),dim=1,index=idx.long()) # 对列索引

print(ans)
# tensor([[108, 107, 109],
#         [101, 103, 102],
#        [103, 101, 106],
#        [107, 108, 109]])