torch.tensor_split¶
- torch.tensor_split(input, indices_or_sections, dim=0) 张量列表 ¶
将张量拆分为多个子张量,所有子张量都是
input
的视图,沿着维度dim
根据indices_or_sections
指定的索引或节数进行拆分。此函数基于 NumPy 的numpy.array_split()
。- 参数
input (张量) – 要拆分的张量
indices_or_sections (张量, int 或 列表 或 元组 的 int) –
如果
indices_or_sections
是一个整数n
或一个值为n
的零维长整型张量,则input
沿着维度dim
拆分为n
个部分。如果input
沿着维度dim
可以被n
整除,则每个部分的大小相同,为input.size(dim) / n
。如果input
不能被n
整除,则前int(input.size(dim) % n)
个部分的大小为int(input.size(dim) / n) + 1
,其余部分的大小为int(input.size(dim) / n)
。如果
indices_or_sections
是一个整数列表或元组,或一个一维长整型张量,则input
沿着维度dim
在列表、元组或张量中的每个索引处进行拆分。例如,indices_or_sections=[2, 3]
和dim=0
将产生张量input[:2]
、input[2:3]
和input[3:]
。如果
indices_or_sections
是一个张量,它必须是 CPU 上的零维或一维长整型张量。dim (int, 可选) – 沿着哪个维度拆分张量。默认值:
0
示例
>>> x = torch.arange(8) >>> torch.tensor_split(x, 3) (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7])) >>> x = torch.arange(7) >>> torch.tensor_split(x, 3) (tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6])) >>> torch.tensor_split(x, (1, 6)) (tensor([0]), tensor([1, 2, 3, 4, 5]), tensor([6])) >>> x = torch.arange(14).reshape(2, 7) >>> x tensor([[ 0, 1, 2, 3, 4, 5, 6], [ 7, 8, 9, 10, 11, 12, 13]]) >>> torch.tensor_split(x, 3, dim=1) (tensor([[0, 1, 2], [7, 8, 9]]), tensor([[ 3, 4], [10, 11]]), tensor([[ 5, 6], [12, 13]])) >>> torch.tensor_split(x, (1, 6), dim=1) (tensor([[0], [7]]), tensor([[ 1, 2, 3, 4, 5], [ 8, 9, 10, 11, 12]]), tensor([[ 6], [13]]))