本文介绍如何在 PyTorch 中仅通过一次张量操作，从高维张量（如 (100, 50, 3)）中按指定行索引和各自独立的起始/长度提取连续子序列，并合并为统一结果，避免显式循环或多次索引拼接。
本文介绍如何在 pytorch 中仅通过一次张量操作，从高维张量（如 `(100, 50, 3)`）中按指定行索引和各自独立的起始/长度提取连续子序列，并合并为统一结果，避免显式循环或多次索引拼接。
在实际深度学习与数据处理中，常需从批量样本（如 batch × seq_len × feature_dim）中对每个样本提取不同起始位置、但相同长度的子序列（例如：第0个样本取 [5:10]，第1个样本取 [10:15]）。若使用传统方式——先索引再逐个切片最后拼接——不仅代码冗余，还无法充分利用 GPU 并行性。PyTorch 提供了更高效、可微分的解决方案：基于 torch.gather 的向量化索引。
核心思路是：将“每行不同起始位置 + 固定长度”的切片需求，转化为一个二维索引矩阵（形状为 [N, L]），再将其广播扩展至目标维度，最后沿指定轴 gather 原张量。
以下为完整实现步骤（以原始张量 arr = torch.randint(0, 9, (100, 50, 3)) 为例）：
import torch # 原始张量：(100, 50, 3) arr = torch.randint(0, 9, (100, 50, 3)) # 目标行索引（第一维）和对应起始位置 row_indices = torch.tensor([5, 55]) # 取第5行和第55行 → shape: [2] start_indices = torch.tensor([5, 10]) # 各自行的起始偏移 → shape: [2] length = 5 # 每行提取长度（必须一致） # 步骤1：构造相对偏移序列 [0, 1, ..., length-1] offsets = torch.arange(length, device=arr.device) # shape: [5] # 步骤2：广播相加，生成绝对列索引矩阵 # start_indices[:, None] + offsets[None, :] → shape: [2, 5] idx_2d = start_indices.unsqueeze(1) + offsets.unsqueeze(0) # 或写作: start_indices[:, None] + offsets[None, :] # 步骤3：扩展至第三维（特征维度），适配 gather 输入要求 # idx_2d.shape = [2, 5] → 扩展为 [2, 5, 3]，使每个 (i,j) 对应全部3个通道 idx_3d = idx_2d.unsqueeze(-1).expand(-1, -1, arr.size(-1)) # shape: [2, 5, 3] # 步骤4：在 dim=1（即序列维度）上 gather —— 注意：arr[row_indices] 后 shape 为 [2, 50, 3] partial_arr = arr[row_indices] # shape: [2, 50, 3] result = partial_arr.gather(1, idx_3d) # shape: [2, 5, 3] # 最终结果：若需展平为 [10, 3]（即 cat([first_result, second_result]) 效果） final_result = result.view(-1, arr.size(-1)) # shape: [10, 3]
✅ 关键要点说明：
- torch.gather 要求 index 张量与 input 在除 dim 外其余维度一致；因此需将二维索引 idx_2d 显式扩展至三维以匹配 partial_arr 的通道数。
- 所有操作均为向量化、无 Python 循环，支持 CUDA 加速与梯度回传（gather 是可微操作）。
- 限制条件：各子序列长度必须相同（因 torch.arange(length) 生成统一 offset 序列）；若长度不等，需 padding 或分组处理，无法单次 gather 完成。
? 进阶提示：若需动态长度，可结合 torch.nn.utils.rnn.pad_sequence 预处理，或使用 torch.scatter 反向构建掩码，但复杂度显著上升。对于绝大多数序列抽取任务（如 NLP 中的 span extraction、时序模型中的 sliding window），固定长度 + gather 是最优实践。
综上，通过 unsqueeze、expand 与 gather 的组合，我们成功将原本需三次独立索引+拼接的操作，压缩为一次全张量运算——兼顾简洁性、性能与可扩展性。
