使用另一个矩阵的Torch查询矩阵

Question

我有一个m张量(张量1)和另一个kx2张量(张量2)，我希望用基于张量2的指数来提取张量1的所有值。

Tensor1
  1   2   3   4   5
  6   7   8   9  10
 11  12  13  14  15
 16  17  18  19  20
[torch.DoubleTensor of size 4x5]

Tensor2
 2  1
 3  5
 1  1
 4  3
[torch.DoubleTensor of size 4x2]

功能就会产生；

6
15
1
18

Answer 1

想到的第一个解决方案是简单地遍历索引并选择相应的值：

function get_elems_simple(tensor, indices)
    local res = torch.Tensor(indices:size(1)):typeAs(tensor)
    local i = 0
    res:apply(
        function () 
            i = i + 1
            return tensor[indices[i]:clone():storage()] 
        end)
    return res
end

这里，tensor[indices[i]:clone():storage()]只是从多维张量中挑选元素的一种通用方法。在k维情况下，这完全类似于tensor[{indices[i][1], ... , indices[i][k]}]。

如果不需要提取大量的值，这种方法工作得很好(瓶颈是:apply方法，它无法使用许多优化技术和SIMD指令，因为它执行的函数是一个黑匣子)。这项工作可以更有效地完成:方法:index完全可以完成您所需要的工作。有一个一维张量。多维目标/索引张量需要扁平：

function flatten_indices(sp_indices, shape)
    sp_indices = sp_indices - 1
    local n_elem, n_dim = sp_indices:size(1), sp_indices:size(2)
    local flat_ind = torch.LongTensor(n_elem):fill(1)

    local mult = 1
    for d = n_dim, 1, -1 do
        flat_ind:add(sp_indices[{{}, d}] * mult)
        mult = mult * shape[d]
    end
    return flat_ind
end

function get_elems_efficient(tensor, sp_indices)
    local flat_indices = flatten_indices(sp_indices, tensor:size()) 
    local flat_tensor = tensor:view(-1)
    return flat_tensor:index(1, flat_indices)
end

差别是很大的：

n = 500000
k = 100
a = torch.rand(n, k)
ind = torch.LongTensor(n, 2)
ind[{{}, 1}]:random(1, n)
ind[{{}, 2}]:random(1, k)

elems1 = get_elems_simple(a, ind)      # 4.53 sec
elems2 = get_elems_efficient(a, ind)   # 0.05 sec

print(torch.all(elems1:eq(elems2)))    # true