如何按指定索引列表分割PyTorch数组？

你可以通过切片张量结合torch.split()方法轻松实现这个需求，步骤如下：

1. 先计算分割长度列表的总和，截取张量的前对应长度部分（因为输入张量长度大于该总和）
2. 使用torch.split()将截取后的张量按指定的长度列表分割成多个子张量

示例代码 #

import torch

# 输入张量
x = torch.arange(20)
# 分割长度列表
splits = [1,2,5,10]

# 截取前sum(splits)个元素
target_tensor = x[:sum(splits)]
# 执行分割
split_result = torch.split(target_tensor, splits)

# 输出结果
print(split_result)

输出结果 #

(tensor([0]),
 tensor([1, 2]),
 tensor([3, 4, 5, 6, 7]),
 tensor([ 8,  9, 10, 11, 12, 13, 14, 15, 16, 17]))

说明 #

• torch.split()方法接受长度列表参数时，会严格按照列表数值将张量分割为对应长度的子张量，要求列表元素之和等于输入张量的长度，因此需要先截取张量到对应长度。
• 这种方法比手动循环切片更高效，符合PyTorch的张量操作习惯。

内容的提问来源于stack exchange，提问作者piccolo