pytorch高阶操作

pytorch高阶操作

where函数

torch.where(condition,x,y)

可能新生成的tensor一部分来自x，一部分来自y，但是是没有规律的

例子：假设一个tensor表示识别概率，大于0.5表示1，小于0.5表示0

a = torch.rand(2,2)
print(a)

tensor([[0.9872, 0.9270],
        [0.6795, 0.0959]])


aa = torch.zeros(2,2)
bb = torch.ones(2,2)

answer = torch.where(a>0.5,aa,bb)
print(answer)

tensor([[0., 0.],
        [0., 1.]])

gather函数

实际就是一个查表的函数

比如像手写数字的识别，【4，10】4张图片，最后识别出每张图片中10个概率最大的index（一般index为几这个数字就是几），但是如果我们的标签不是1~10，而是另外有一张表来对应，不同的index对应不同的标签，这时就可以使用gather函数

例子：

prob = torch.rand(4,10)

idx = prob.topk(3,dim=1)
idx1 = idx[1]

print(idx1)

tensor([[1, 3, 4],
        [2, 0, 3],
        [5, 4, 2],
        [9, 4, 5]])

label = torch.arange(10)+100#为了方面随便初始化的label

print(torch.gather(label.expand(4,10),dim=1,index=idx1.long()))


tensor([[101, 103, 104],
        [102, 100, 103],
        [105, 104, 102],
        [109, 104, 105]])