pytorch对数值排序问题

1）pytorch选出数据中的最大值索引

torch.argmax()函数中dim的不同值表示在该维度下找出tensor的最大值的索引。

b=torch.argmax(a,dim=1)#输出a在第一维度的最大值的索引

2）pytorch选出数据中的前k个最大（最小）值及其索引

data_sorted, idx_sorted = torch.sort(data, descending=True)
#descending为alse，升序，为True，降序
idx = idx_sorted[:k]

data_sorted:排好序的数据
idx_sorted：对应排序数据的索引

因此只需设置k的大小，就可以截取到前k个最大值的索引,但是缺点是打乱了原来数据的顺序，返回的是重新排序的数据，如果是处理两个具有相同含义索引的数据，则不能用这种方法。

torch.topk 

import torch

item=torch.IntTensor([1,2,4,7,3,2])
value,indices=torch.topk(item,3)
print("value:",value)
print("indices:",indices)

#value中存储的是对应的top3的元素，并按照从大到小的取值方式进行存储

#indices中存储的是value中top3元素在原Tensor中的索引值

Pytorch 替换tensor中大于某个值的所有元素torch.where

a = torch.rand((1, 3))

# 保证和a相同的维度大小
zero = torch.zeros_like(a)
one = torch.ones_like(a)

# a中大于0.5的用one(1)替换,否则a替换,即不变
a = torch.where(a > 0.5, one, a)
 
# a中小于0.5的用zero(0)替换,否则a替换,即不变
a = torch.where(a < 0.5, zero, a)

显然 torch.where及查找阈值与替换为一身，这个用起来比较方便。