pytorch 基本操作（四）——进阶

这节简单介绍两个并行的函数。

where

where函数相当于excel里的if函数如过满足条件，则选择前项，否则后项。

import torch
judge = torch.rand(2,2)
judge
'''
tensor([[0.1112, 0.5292],
        [0.1838, 0.5912]])
'''

a = torch.zeros(2,2)
b = torch.ones(2,2)
torch.where(judge>0.5, a, b)
'''
tensor([[1., 0.],
        [1., 0.]])
'''

gather

gather函数是用来做数据查询的，首先建立一个3x4的矩阵然后建一个索引。

norm = torch.randn(3,4)
ind = norm.topk(k = 2)[1]
ind
'''
tensor([[2, 3],
        [0, 1],
        [1, 0]])
'''

接着利用gather来索引特定位置的值。

ex_gather = torch.arange(4)+10
ex_gather.expand(3,4)
'''
tensor([[10, 11, 12, 13],
        [10, 11, 12, 13],
        [10, 11, 12, 13]])
'''
torch.gather(ex_gather.expand(3,4), dim = 1, index = ind)
'''
tensor([[12, 13],
        [10, 11],
        [11, 10]])
'''

下一节，就要进入pytorch最为重要的梯度概念了。

Ref

[1] https://github.com/irobbwu/pytorch-intro/blob/main/03.advanced%20op.ipynb