如何高效地从二维张量的每一行中提取不同起始位置的固定长度子序列_网络技术

如何高效地从二维张量的每一行中提取不同起始位置的固定长度子序列

发布时间 - 2026-01-06 00:00:00 点击率：次

本文介绍一种无需显式循环即可从 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
该方法天然支持 GPU 张量，全程无 CPU-GPU 数据拷贝，性能远超循环；
可轻松扩展至更高维（如 (B, T, D) 中按 B 维切时间步），只需调整 dim 和索引广播逻辑。

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

# python # ai # 深度学习 # pytorch

相关栏目：【网站优化151355 】【网络推广146373 】【网络技术251813 】【 AI营销90571 】

上一篇：把编译安装的httpd，利用service和chkconfi

下一篇：长佩文学网页入口长佩文学官方网站入口点此

相关栏目网站优化
网络推广
网络技术
AI营销

最新文章 Sublime怎么一键压缩JS代码 Su sublime如何在搜索中使用正则表达式 Sublime如何设置透明窗口效果 Su mysql如何设计商品表结构_mysql css属性背景图不显示怎么办_通过检查路如何使用Golang实现排序_Golan 农历闰月是怎么回事_为合回归年加一月调整塑造《刺客信条》艾吉奥传奇的编剧离开育碧 1英里等于多少公里 1mile和km的换 css grid布局中行和列是如何定义的 PS批量旋转和翻转图片，快速校正图片方向 C# Swagger UI自定义方法 C OPPO手机九宫格和全键盘怎么切换_OP Go语言如何实现用户登录注册_Golan 1节飞行速度多少公里每小时 1节是多少公纸嫁衣8千子树第五章庙门怎么开启庙门 Laravel 多行数据编辑表单中实现逐明日之后如何提升钓鱼等级明日之后提升钓支付宝怎样查年度账单_支付宝年度账单查看 C# 多线程UI更新Dispatcher

上一篇：把编译安装的httpd，利用service和chkconfi

下一篇：长佩文学网页入口长佩文学官方网站入口点此