官方文档:https://pytorch.org/docs/1.10.1/generated/torch.split.html#torch.split
torch.split()的作用是把一个tensor拆分为多个tensor,相当于是concat的逆过程,定义如下:
torch.split(tensor, split_size_or_sections, dim=0)
- tensor:待被拆分的tensor
- split_size_or_sections:记该参数的值为k。如果k为整数(int),那么切分得到的每个子tensor在被拆分维度上的长度均为k(不要求严格整除,如果不能整除,最后一个子tensor会小一点);如果k为列表(list),那么第i个子tensor在被拆分维度上的长度为list[i]。
- dim:沿着哪一维度开始拆分
这里举一个简单的例子。比如在计算机视觉任务中,需要处理一个batch size为4,通道数为128,长宽均为512的feature map,按照[B, C, H, W]的格式,创建如下:
import torch
t = torch.rand([4, 128, 512, 512])
如果我们要按通道将其等分为两半,写法如下:
import torch
t = torch.rand([4, 128, 512, 512])
print(t.shape)
a, b = torch.split(t, 64, dim=1)
print(a.shape)
print(b.shape)
输出结果如下:
torch.Size([4, 128, 512, 512])
torch.Size([4, 64, 512, 512])
torch.Size([4, 64, 512, 512])
这里由于是等分,因此可以直接用int指定。而如果是更灵活的切分,则应该用列表,比如我们希望第1个子tensor通道数为3,而第二个子tensor通道数为125,则可以改成:
a, b = torch.split(t, [3, 125], dim=1)
此时输出结果如下:
torch.Size([4, 128, 512, 512])
torch.Size([4, 3, 512, 512])
torch.Size([4, 125, 512, 512])
需要注意的是,如果传入list进去拆分的话,那么每个子tensor的维度加起来应该是要等于原来的维度的。在这个例子中,有3+125=128。如果我们改成3和124:
a, b = torch.split(t, [3, 124], dim=1)
这种情况下是会报错的。