torch.split¶
- torch.split(tensor, split_size_or_sections, dim=0)[源代码]¶
将张量分割成块。每个块都是原始张量的视图。
如果
split_size_or_sections
是整数类型,则tensor
将被分割成大小相等的块(如果可能)。如果张量沿给定维度dim
的大小不能被split_size
整除,则最后一个块将更小。如果
split_size_or_sections
是列表,则tensor
将被分割成len(split_size_or_sections)
个块,其在dim
维度上的大小根据split_size_or_sections
确定。- 参数
- 返回类型
示例
>>> a = torch.arange(10).reshape(5, 2) >>> a tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]) >>> torch.split(a, 2) (tensor([[0, 1], [2, 3]]), tensor([[4, 5], [6, 7]]), tensor([[8, 9]])) >>> torch.split(a, [1, 4]) (tensor([[0, 1]]), tensor([[2, 3], [4, 5], [6, 7], [8, 9]]))