Pytorch torch.split()的简单用法

官方文档：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)
1

• 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])
1
2

如果我们要按通道将其等分为两半，写法如下：

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)
1
2
3
4
5
6

输出结果如下：

torch.Size([4, 128, 512, 512])
torch.Size([4, 64, 512, 512])
torch.Size([4, 64, 512, 512])
1
2
3

这里由于是等分，因此可以直接用int指定。而如果是更灵活的切分，则应该用列表，比如我们希望第1个子tensor通道数为3，而第二个子tensor通道数为125，则可以改成：

a, b = torch.split(t, [3, 125], dim=1)
1

此时输出结果如下：

torch.Size([4, 128, 512, 512])
torch.Size([4, 3, 512, 512])
torch.Size([4, 125, 512, 512])
1
2
3

需要注意的是，如果传入list进去拆分的话，那么每个子tensor的维度加起来应该是要等于原来的维度的。在这个例子中，有3+125=128。如果我们改成3和124：

a, b = torch.split(t, [3, 124], dim=1)
1

这种情况下是会报错的。