本文介绍如何在 pytorch 中仅通过一次张量操作，高效提取多个样本（沿第一维）各自指定范围的连续子序列，避免多次索引与拼接，提升代码简洁性与运行效率。
本文介绍如何在 pytorch 中仅通过一次张量操作，高效提取多个样本（沿第一维）各自指定范围的连续子序列，避免多次索引与拼接，提升代码简洁性与运行效率。
在处理高维张量（如 (N, L, D) 形状的序列数据）时，常需对每个样本（即第 0 维的每个元素）独立提取不同起始位置、相同长度的子序列。例如，从 arr = torch.randint(0, 9, (100, 50, 3)) 中：
- 对第 5 个样本（索引 5）取 5:10（共 5 个时间步）；
- 对第 55 个样本（索引 55）取 10:15（同样 5 个时间步）。
若逐个切片再拼接（torch.cat），不仅冗余，还破坏了向量化优势。正确解法是利用 广播 + torch.gather 实现单次批量索引。
✅ 核心步骤详解
-
构建批量索引矩阵
设目标样本索引为 row_indices = [5, 55]，各自行内起始位置为 starts = [5, 10]，长度统一为 L = 5。
先生成相对偏移：torch.arange(L) → [0,1,2,3,4]；
再通过广播加法构造绝对位置矩阵：starts = torch.tensor([5, 10]) offsets = torch.arange(5) # shape: (5,) idx_ = starts.unsqueeze(1) + offsets.unsqueeze(0) # shape: (2, 5) # tensor([[ 5, 6, 7, 8, 9], # [10, 11, 12, 13, 14]])
-
扩展维度以匹配目标张量
原张量 arr 的最后一维是特征维度（如 D=3），需将 idx_ 扩展为 (2, 5, 3)，使其可被 gather 沿 dim=1 使用：idx_expanded = idx_.unsqueeze(-1).expand(-1, -1, arr.size(-1)) # shape: (2, 5, 3)
-
执行批量 gather
先用 arr[row_indices] 获取目标样本子集（partial_arr），再沿 dim=1（时间维）gather：row_indices = torch.tensor([5, 55]) partial_arr = arr[row_indices] # shape: (2, 50, 3) result = partial_arr.gather(1, idx_expanded) # shape: (2, 5, 3)
此时 result[i] 即为第 row_indices[i] 个样本中从 starts[i] 开始的 L 个元素。
⚠️ 注意事项：
- torch.gather 要求 index 张量与输入张量除 dim 外其余维度一致；此处 dim=1，故 idx_expanded 需匹配 (2, 5, 3)。
- 若各切片长度不等，无法直接广播，需改用 torch.stack + 列表推导式（牺牲向量化），或使用 torch.nn.utils.rnn.pad_sequence 预处理。
- 此方法天然支持 GPU 加速，且内存友好——无需中间 cat 拼接。
最终结果为形状 (2, 5, 3) 的张量，每行对应一个定制化切片，真正实现「一次操作，批量提取」。
