给定 Torch ( torch.Tensor
) 中的一维张量,包含可以比较的值(例如浮点),我们如何提取该张量中前k个值的索引?
除了蛮力方法,我正在寻找一些 API 调用,由 Torch/lua 提供,它可以有效地执行这个任务。
从拉取请求#496开始, Torch 现在包含一个名为torch.topk
. 例子:
> t = torch.Tensor{9, 1, 8, 2, 7, 3, 6, 4, 5}
-- obtain the 3 smallest elements
> res = t:topk(3)
> print(res)
1
2
3
[torch.DoubleTensor of size 3]
-- you can also get the indices in addition
> res, ind = t:topk(3)
> print(ind)
2
4
6
[torch.LongTensor of size 3]
-- alternatively you can obtain the k largest elements as follow
-- (see the API documentation for more details)
> res = t:topk(3, true)
> print(res)
9
8
7
[torch.DoubleTensor of size 3]
在撰写本文时,CPU 实现遵循了一种严格的方法(有计划在未来对其进行改进)。话虽如此,目前正在审查针对 cutorch 的优化 GPU 实现。
您可以使用topk功能。
例如:
import torch
t = torch.tensor([5.7, 1.4, 9.5, 1.6, 6.1, 4.3])
values,indices = t.topk(2)
print(values)
print(indices)
结果:
tensor([9.5000, 6.1000])
tensor([2, 4])
只需遍历张量并运行您的比较:
require 'torch'
data = torch.Tensor({1,2,3,4,505,6,7,8,9,10,11,12})
idx = 1
max = data[1]
for i=1,data:size()[1] do
if data[i]>max then
max=data[i]
idx=i
end
end
print(idx,max)
--EDIT-- 响应您的编辑:使用此处记录的 torch.max 操作:https ://github.com/torch/torch7/blob/master/doc/maths.md#torchmaxresval-resind-x-dim .. .
y, i = torch.max(x, 1) returns the largest element in each column (across rows) of x, and a Tensor i of their corresponding indices in x