二维张量固定行子序列提取方法

小伙伴们对文章编程感兴趣吗？是否正在学习相关知识点？如果是，那么本文《二维张量每行提取固定长度子序列方法》，就很适合你，本篇文章讲解的知识点主要包括。在之后的文章中也会多多分享相关知识点，希望对大家的知识积累有所帮助！

本文介绍一种无需显式循环即可从 PyTorch 二维张量各行中按指定起始索引和统一长度批量切片的方法，核心是利用 `torch.arange` 构造索引张量，并通过 `gather` 实现向量化索引选取。

在深度学习与科学计算中，常需对批量数据（如序列、特征图）进行行级动态切片——即每行按各自起始位置截取固定长度的子序列。若使用 Python 循环或列表推导式逐行处理，不仅代码冗长，还会严重损害 GPU 加速优势。幸运的是，PyTorch 提供了完全向量化的解决方案：torch.gather + 动态索引张量构造。

其关键思想是：将每行的切片需求（起始索引 + 长度）转化为一个与输出形状一致的列索引张量，再通过 gather(dim=1, index=...) 沿列维度（dim=1）收集对应元素。

以下是完整实现：

import torch

def batch_row_slice(data: torch.Tensor, start_indices: torch.Tensor, length: int, dim: int = 1) -> torch.Tensor:
    """
    从2D张量每行中提取固定长度子序列（向量化实现）

    Args:
        data: 输入2D张量，shape = (N, D)
        start_indices: 每行起始列索引，shape = (N,)，dtype=torch.long
        length: 每个切片的固定长度（所有行相同）
        dim: 沿哪个维度切片（默认为列维度，即 dim=1）

    Returns:
        切片结果，shape = (N, length)
    """
    if dim != 1:
        raise NotImplementedError("仅支持沿列维度（dim=1）切片")

    # 构造每行的索引序列：[[s0, s0+1], [s1, s1+1], ..., [s_{N-1}, s_{N-1}+1]]
    # 使用 broadcasting 和 arange 实现高效生成
    arange = torch.arange(length, device=data.device, dtype=start_indices.dtype)
    # start_indices: (N,) → (N, 1); arange: (length,) → (1, length)
    # 广播后得到 (N, length) 的起始偏移矩阵
    index_tensor = start_indices.unsqueeze(1) + arange.unsqueeze(0)

    # 使用 gather 沿 dim=1 收集对应列元素
    return torch.gather(data, dim=1, index=index_tensor)

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

start_idx = torch.tensor([0, 3, 1], dtype=torch.long)  # 注意：必须为 long 类型！
result = batch_row_slice(data, start_idx, length=2)

print(result)
# 输出：
# tensor([[ 1.,  2.],
#         [ 9., 10.],
#         [12., 13.]])

✅ 关键要点说明：

• start_indices 必须是 torch.long 类型（浮点型索引不被 gather 接受）；
• length 必须对所有行保持一致，这是向量化前提；
• 索引越界行为由 gather 自动处理（超出范围时返回 0 填充），生产环境建议提前校验 start_idx + length <= data.size(1)；
• 该方法天然支持 GPU 张量，全程无 CPU-GPU 数据拷贝，性能远超循环；
• 可轻松扩展至更高维（如 (B, T, D) 中按 B 维切时间步），只需调整 dim 和索引广播逻辑。

通过此方法，你可以在保持代码简洁性的同时，充分发挥 PyTorch 的张量并行能力，适用于 Transformer attention mask 构建、时序窗口采样、动态 padding 等典型场景。

以上就是《二维张量固定行子序列提取方法》的详细内容，更多关于的资料请关注golang学习网公众号！