Pytorch: split() 函数

作用：对tensor在某一dim维度下，根据指定的大小split_size=int，或者list(int)来分割数据，返回tuple元组。

torch.split(tensor, split_size_or_sections, dim=0)

注意两个细节：1、当split_size为一个int数时，若不能整除int，剩余的数据直接作为一块。2、当split_size为一个列表时，所有数字的和等于要分割的维度大小。

#细节1
import torch
input = torch.randn(4,10,8)
output = torch.split(input, 3, dim=1) #3+3+3+1=10,不够的单独作为一块
print(type(output))
for chunk in output:
    print(chunk.size)
>>>
<class 'tuple'>
torch.Size([4, 3, 8])
torch.Size([4, 3, 8])
torch.Size([4, 3, 8])
torch.Size([4, 1, 8]) 

#细节2
import torch
input = torch.randn(4,10,8)
split_list = [3, 3, 3, 1]
output = torch.split(input, split_list, dim=1) #3+3+3+1=10
print(type(output))
for chunk in output:
    print(chunk.size())
>>>
<class 'tuple'>
torch.Size([4, 3, 8])
torch.Size([4, 3, 8])
torch.Size([4, 3, 8])
torch.Size([4, 1, 8])