本文介绍如何使用 torch.tensor_split() 结合累积和，将 pytorch 张量按指定长度列表（而非切片位置）精准分割为多个子张量，并提供可直接运行的代码示例与关键注意事项。
本文介绍如何使用 torch.tensor_split() 结合累积和，将 pytorch 张量按指定长度列表（而非切片位置）精准分割为多个子张量，并提供可直接运行的代码示例与关键注意事项。
在 PyTorch 实际开发中，常需根据一组「段长度」（如 [1, 2, 5, 10]）将一维张量划分为若干连续子张量，而非基于绝对索引位置（如 [1, 3, 8, 18]）进行切分。此时，torch.split() 虽然支持按长度列表分割，但要求所有段长度之和必须严格等于输入张量长度；而 torch.tensor_split() 更灵活，允许按“分割点”（split points）划分，且不强制覆盖全部元素——这使其成为本场景的理想选择。
核心思路是：将长度列表转换为累积切分位置（cumulative split indices），即计算前缀和，作为 tensor_split 的切分依据。注意，tensor_split 的 split_size_or_sections 参数在此处应传入切分点列表（如 [1, 3, 8, 18]），它表示在这些索引之前进行切割（左闭右开），最终会生成 len(split_points) + 1 个片段。由于题目示例仅需前 len(splits) 段（最后一段被舍弃），需手动截取 [:-1]。
以下是完整、健壮的实现代码：
import torch
import numpy as np
# 示例数据
t = torch.arange(20) # shape: (20,)
splits = [1, 2, 5, 10] # 各段期望长度
# 计算累积切分点：[1, 1+2, 1+2+5, 1+2+5+10] → [1, 3, 8, 18]
split_points = np.cumsum(splits).tolist()
# 使用 tensor_split 切分（返回 len(split_points)+1 个张量）
chunks = t.tensor_split(split_points)
# 取前 len(splits) 个片段（舍弃最后一段，因题目示例未包含剩余部分）
result = chunks[:-1]
for i, chunk in enumerate(result):
print(f"Chunk {i}: {chunk}")
输出结果完全匹配预期：
Chunk 0: tensor([0]) Chunk 1: tensor([1, 2]) Chunk 2: tensor([3, 4, 5, 6, 7]) Chunk 3: tensor([ 8, 9, 10, 11, 12, 13, 14, 15, 16, 17])
⚠️ 注意事项：
- torch.tensor_split() 自 PyTorch 1.8+ 起可用，低版本请升级或改用 torch.narrow() 手动实现；
- np.cumsum() 可替换为纯 PyTorch 实现（如 torch.tensor(splits).cumsum(0).tolist()），避免 NumPy 依赖；
- 若 sum(splits) >= len(t)，最后一段可能为空或越界，建议提前校验：assert sum(splits)
- tensor_split 返回的是 tuple，如需 list 可显式调用 list(chunks[:-1])。
总结：该方法兼顾简洁性与鲁棒性，适用于数据预处理、批量采样、序列分段等任务，是 torch.split() 的有力补充。
