Pytorch文档学习 TORCH.TENSOR_SPLIT

torch.tensor_split

torch.tensor_split(input, indices_or_sections, dim=0) → List of Tensors

这个方法和split方法长得很像。他们的作用都是根据indices_or_sections，把输入拆分成几个视图。区别在于：

①split的indices_or_sections的类型为int或者list，而tensor_split的类型为tensor，int，list，tuple of list。

②当indices_or_sections类型为int时，对于split，其按照int的值n不断划分，直到最后一个取不完分为最后一组。对于tensor_split，会优先计算第一个分组，大小为int(size/n) +1，然后剩下的分组大小尽可能一样。

③当indices_or_sections类型为list时，对于split，划分值之和不得超过元素个数，比如不能写成【3,4,15】，否则会报错。对于tensor_split，其划分的方法是对list或者tuple中的元素，根据元素组成的区间进行划分。如对【3,4,15】，下标小于3的，下标小于4大于3的，下标小于6大于4的，下标小于15大于6的，该方法会划分出以上4组。下标超出实际大小的时候会返回一个空的tensor。

x = torch.arange(7)
print(torch.tensor_split(x, 3))
print(torch.split(x, 3))

(tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6]))
(tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6]))

x = torch.arange(7)
print(torch.tensor_split(x, (3,4,15)))
print(torch.split(x, [3,4]))

(tensor([0, 1, 2]), tensor([3]), tensor([4, 5, 6]), tensor([], dtype=torch.int64))
(tensor([0, 1, 2]), tensor([3, 4, 5, 6]))