作用:对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])
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐