从1-D Tensor中提取top-k值指数

时间:2016-01-12 17:36:25

标签: lua torch

给定Torch中的1-D张量(torch.Tensor),包含可以比较的值(比如浮点),我们如何提取top- k 值的索引张量?

除了暴力方法之外,我正在寻找一些可以有效执行此任务的API调用,即Torch / lua提供的。

3 个答案:

答案 0 :(得分:4)

截取请求#496 Torch现在包含一个名为torch.topk的内置API。例如:

> 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实现遵循sort and narrow approach(有计划在未来对其进行改进)。据说,针对cutorch的优化GPU实现目前正在reviewed

答案 1 :(得分:0)

只需遍历张量并运行比较:

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

答案 2 :(得分:0)

您可以使用topk功能。

例如:

def I(i: Int): AnyRef

结果:

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)