indexing - 张量的 torch 逻辑索引

Question

我正在寻找一种优雅的方法来选择满足某些约束的火炬张量的子集。
例如，假设我有:

A = torch.rand(10,2)-1

和 S是一个 10x1 张量，

sel = torch.ge(S,5) -- this is a ByteTensor

我希望能够进行逻辑索引，如下所示:

A1 = A[sel]

但这不起作用。
所以有 index接受 LongTensor 的函数但我找不到一种简单的方法来转换 S到 LongTensor ，以下情况除外:

sel = torch.nonzero(sel)

它返回一个 K x 2 张量(K 是 S >= 5 的值的数量)。那么我必须将其转换为一维数组，这最终允许我对 A 进行索引:

A:index(1,torch.squeeze(sel:select(2,1)))

这很麻烦；在例如Matlab 我所要做的就是

A(S>=5,:)

任何人都可以提出更好的方法吗？

Answer 1

一种可能的选择是:

sel = S:ge(5):expandAs(A)   -- now you can use this mask with the [] operator
A1 = A[sel]:unfold(1, 2, 2) -- unfold to get back a 2D tensor

例子:

> A = torch.rand(3,2)-1
-0.0047 -0.7976
-0.2653 -0.4582
-0.9713 -0.9660
[torch.DoubleTensor of size 3x2]

> S = torch.Tensor{{6}, {1}, {5}}
 6
 1
 5
[torch.DoubleTensor of size 3x1]

> sel = S:ge(5):expandAs(A)
1  1
0  0
1  1
[torch.ByteTensor of size 3x2]

> A[sel]
-0.0047
-0.7976
-0.9713
-0.9660
[torch.DoubleTensor of size 4]

> A[sel]:unfold(1, 2, 2)
-0.0047 -0.7976
-0.9713 -0.9660
[torch.DoubleTensor of size 2x2]