如何无需迭代从2D Tensor的每行选取指定范围的元素？

问题 #

假设有一个2D张量：

data = tensor([[ 1.,  2.,  3.,  4.,  5.],
               [ 6.,  7.,  8.,  9., 10.],
               [11., 12., 13., 14., 15.]])

另有两个1D张量作为每行待选取元素的起始和结束索引：

start = tensor([0., 3., 1.])
end = start + 2               # end始终比start大2

有没有不用迭代data的方法，直接选取第i行中start[i]:end[i]范围内的元素？

示例预期输出：

tensor([[ 1.,  2.],
        [ 9., 10.],
        [12., 13.]])

解决方案 #

当然可以，不用循环迭代就能实现，核心是构造每行对应的列索引矩阵，再结合行索引提取元素：

1. 先把start转成整数类型（索引必须是整数）：

   start = start.long()
   

2. 构造列偏移量：因为每个区间固定长度为2，偏移量取[0,1]，并广播成和start匹配的形状：

   offsets = torch.arange(2, device=data.device)  # 生成tensor([0,1])
   col_indices = start.unsqueeze(1) + offsets  # 得到形状为(3,2)的列索引矩阵
   

3. 构造行索引：每行对应自身的索引，同样广播成(3,2)的形状：

   row_indices = torch.arange(data.shape[0], device=data.device).unsqueeze(1).repeat(1,2)
   

4. 提取元素，两种方法任选：
   • 方法一：用torch.gather

     result = data.gather(1, col_indices)
     

   • 方法二：直接用双索引取值

     result = data[row_indices, col_indices]
     

两种方法都能直接得到预期输出，全程不需要逐行迭代。

内容的提问来源于stack exchange，提问作者Belleroph0N