torch.split()

torch.split()

官网链接：https://pytorch.org/docs/stable/torch.html
官网解释：Splits the tensor into chunks.——PyTorch中用于分割张量的函数。
作用：将一个多维张量分割成多个张量。

If split_size_or_sections is an integer type, then tensor will be split into equally sized chunks (if possible). Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.
If split_size_or_sections is a list, then tensor will be split into len(split_size_or_sections) chunks with sizes in dim according to split_size_or_sections.

定义是:

torch.split(tensor, split_size_or_sections, dim=0)
参数解释:
- tensor:要分割的输入张量
- split_size_or_sections:
    - 如果是一个整数,则表示分割成每个张量里包含split_size_or_sections个张量，而不是分成split_size_or_sections个
    - 如果是一个列表,则表示对dim维度进行分割，分割为指定大小的张量
- dim:沿着哪个维度进行分割,默认是dim=0，第一维

例1,有这样一个3D张量:

# 生成大小为(2, 4, 8)的随机张量
random_tensor = torch.rand(2, 4, 8)
tensor([[[0.8644, 0.0177, 0.7970, 0.7016, 0.4632, 0.9147, 0.8053, 0.4261],
         [0.0450, 0.9565, 0.8375, 0.9347, 0.8196, 0.8751, 0.4523, 0.9660],
         [0.8350, 0.1566, 0.8367, 0.0345, 0.6804, 0.7308, 0.8989, 0.9943],
         [0.2294, 0.2361, 0.1537, 0.9923, 0.7680, 0.0824, 0.3566, 0.6546]],

        [[0.2106, 0.2736, 0.8687, 0.4333, 0.4102, 0.4820, 0.7104, 0.7776],
         [0.6558, 0.1098, 0.4384, 0.4891, 0.3681, 0.7371, 0.2555, 0.2687],
         [0.4181, 0.6644, 0.1816, 0.2111, 0.8317, 0.4180, 0.7011, 0.7221],
         [0.1922, 0.4405, 0.6633, 0.5787, 0.9912, 0.0370, 0.9894, 0.8748]]])

我们可以这样分割:

1. torch.split(random_tensor, 2, dim=1) ：分割第二维(dim=1)

split_2 = torch.split(random_tensor, 2, dim=1)  # 返回一个元组 tuple
split_2
# (tensor([[[0.8644, 0.0177, 0.7970, 0.7016, 0.4632, 0.9147, 0.8053, 0.4261],
#           [0.0450, 0.9565, 0.8375, 0.9347, 0.8196, 0.8751, 0.4523, 0.9660]],
 
#          [[0.2106, 0.2736, 0.8687, 0.4333, 0.4102, 0.4820, 0.7104, 0.7776],
#           [0.6558, 0.1098, 0.4384, 0.4891, 0.3681, 0.7371, 0.2555, 0.2687]]]),
#  tensor([[[0.8350, 0.1566, 0.8367, 0.0345, 0.6804, 0.7308, 0.8989, 0.9943],
#           [0.2294, 0.2361, 0.1537, 0.9923, 0.7680, 0.0824, 0.3566, 0.6546]],
 
#          [[0.4181, 0.6644, 0.1816, 0.2111, 0.8317, 0.4180, 0.7011, 0.7221],
#           [0.1922, 0.4405, 0.6633, 0.5787, 0.9912, 0.0370, 0.9894, 0.8748]]]))
len(split_2)   # 2   
split_2[0]
# tensor([[[0.8644, 0.0177, 0.7970, 0.7016, 0.4632, 0.9147, 0.8053, 0.4261],
#          [0.0450, 0.9565, 0.8375, 0.9347, 0.8196, 0.8751, 0.4523, 0.9660]],

#         [[0.2106, 0.2736, 0.8687, 0.4333, 0.4102, 0.4820, 0.7104, 0.7776],
#          [0.6558, 0.1098, 0.4384, 0.4891, 0.3681, 0.7371, 0.2555, 0.2687]]])
split_2[0].size()  # torch.Size([2, 2, 8])
split_2[1]
# tensor([[[0.8350, 0.1566, 0.8367, 0.0345, 0.6804, 0.7308, 0.8989, 0.9943],
#          [0.2294, 0.2361, 0.1537, 0.9923, 0.7680, 0.0824, 0.3566, 0.6546]],

#         [[0.4181, 0.6644, 0.1816, 0.2111, 0.8317, 0.4180, 0.7011, 0.7221],
#          [0.1922, 0.4405, 0.6633, 0.5787, 0.9912, 0.0370, 0.9894, 0.8748]]])
split_2[1].size()  # torch.Size([2, 2, 8])

2. torch.split(random_tensor, 3, dim=1) 与上例对比

split_3 = torch.split(random_tensor, 3, dim=1)
split_3
# (tensor([[[0.8644, 0.0177, 0.7970, 0.7016, 0.4632, 0.9147, 0.8053, 0.4261],
#           [0.0450, 0.9565, 0.8375, 0.9347, 0.8196, 0.8751, 0.4523, 0.9660],
#           [0.8350, 0.1566, 0.8367, 0.0345, 0.6804, 0.7308, 0.8989, 0.9943]],
 
#          [[0.2106, 0.2736, 0.8687, 0.4333, 0.4102, 0.4820, 0.7104, 0.7776],
#           [0.6558, 0.1098, 0.4384, 0.4891, 0.3681, 0.7371, 0.2555, 0.2687],
#           [0.4181, 0.6644, 0.1816, 0.2111, 0.8317, 0.4180, 0.7011, 0.7221]]]),
#  tensor([[[0.2294, 0.2361, 0.1537, 0.9923, 0.7680, 0.0824, 0.3566, 0.6546]],
 
#          [[0.1922, 0.4405, 0.6633, 0.5787, 0.9912, 0.0370, 0.9894, 0.8748]]]))
len(split_3)  # 2   1维长度为4，第一次取3，第二次也应取3，但是剩余长度不够，所以取1
split_3[0]    # torch.Size([2, 3, 8])
# tensor([[[0.8644, 0.0177, 0.7970, 0.7016, 0.4632, 0.9147, 0.8053, 0.4261],
#          [0.0450, 0.9565, 0.8375, 0.9347, 0.8196, 0.8751, 0.4523, 0.9660],
#          [0.8350, 0.1566, 0.8367, 0.0345, 0.6804, 0.7308, 0.8989, 0.9943]],

#         [[0.2106, 0.2736, 0.8687, 0.4333, 0.4102, 0.4820, 0.7104, 0.7776],
#          [0.6558, 0.1098, 0.4384, 0.4891, 0.3681, 0.7371, 0.2555, 0.2687],
#          [0.4181, 0.6644, 0.1816, 0.2111, 0.8317, 0.4180, 0.7011, 0.7221]]])
split_3[1]   # torch.Size([2, 1, 8])
# tensor([[[0.2294, 0.2361, 0.1537, 0.9923, 0.7680, 0.0824, 0.3566, 0.6546]],

#         [[0.1922, 0.4405, 0.6633, 0.5787, 0.9912, 0.0370, 0.9894, 0.8748]]])

3. torch.split(random_tensor, [1, 3], dim=1)

split_1_3 = torch.split(random_tensor, [1, 3], dim=1) # 列表中数值总和必须与原维度数值相等
split_1_3
# (tensor([[[0.8644, 0.0177, 0.7970, 0.7016, 0.4632, 0.9147, 0.8053, 0.4261]],
 
#          [[0.2106, 0.2736, 0.8687, 0.4333, 0.4102, 0.4820, 0.7104, 0.7776]]]),
#  tensor([[[0.0450, 0.9565, 0.8375, 0.9347, 0.8196, 0.8751, 0.4523, 0.9660],
#           [0.8350, 0.1566, 0.8367, 0.0345, 0.6804, 0.7308, 0.8989, 0.9943],
#           [0.2294, 0.2361, 0.1537, 0.9923, 0.7680, 0.0824, 0.3566, 0.6546]],
 
#          [[0.6558, 0.1098, 0.4384, 0.4891, 0.3681, 0.7371, 0.2555, 0.2687],
#           [0.4181, 0.6644, 0.1816, 0.2111, 0.8317, 0.4180, 0.7011, 0.7221],
#           [0.1922, 0.4405, 0.6633, 0.5787, 0.9912, 0.0370, 0.9894, 0.8748]]]))
len(split_1_3)  #2
split_1_3[0]  # torch.Size([2, 1, 8])
# tensor([[[0.8644, 0.0177, 0.7970, 0.7016, 0.4632, 0.9147, 0.8053, 0.4261]],

#         [[0.2106, 0.2736, 0.8687, 0.4333, 0.4102, 0.4820, 0.7104, 0.7776]]])
split_1_3[1]   # torch.Size([2, 3, 8])
# tensor([[[0.0450, 0.9565, 0.8375, 0.9347, 0.8196, 0.8751, 0.4523, 0.9660],
#          [0.8350, 0.1566, 0.8367, 0.0345, 0.6804, 0.7308, 0.8989, 0.9943],
#          [0.2294, 0.2361, 0.1537, 0.9923, 0.7680, 0.0824, 0.3566, 0.6546]],

#         [[0.6558, 0.1098, 0.4384, 0.4891, 0.3681, 0.7371, 0.2555, 0.2687],
#          [0.4181, 0.6644, 0.1816, 0.2111, 0.8317, 0.4180, 0.7011, 0.7221],
#          [0.1922, 0.4405, 0.6633, 0.5787, 0.9912, 0.0370, 0.9894, 0.8748]]])

例2,有这样一个3D张量:

random_tensor = torch.rand(2, 2, 3)
tensor([[[0.0445, 0.0481, 0.1199],
         [0.2850, 0.1215, 0.0584]],

        [[0.1323, 0.4458, 0.0899],
         [0.3338, 0.3624, 0.7511]]])

1. torch.split(random_tensor, 2, dim=1)：分割第二维(dim=1)，第一次取两个张量，数据取完。这里本身就是两个张量，所以还是返回自身

split_2 = torch.split(random_tensor, 2, dim=1)   # 返回元组
split_2
# (tensor([[[0.0445, 0.0481, 0.1199],
#           [0.2850, 0.1215, 0.0584]],
 
#          [[0.1323, 0.4458, 0.0899],
#           [0.3338, 0.3624, 0.7511]]]),)
len(split_2)   # 1
split_2[0]     # torch.Size([2, 2, 3])
# tensor([[[0.0445, 0.0481, 0.1199],
#          [0.2850, 0.1215, 0.0584]],

#         [[0.1323, 0.4458, 0.0899],
#          [0.3338, 0.3624, 0.7511]]])
split_2[1]   # 报错

2. torch.split(random_tensor, [1, 2], dim=2)：沿第三维(dim=2)分割

split_1_2 = torch.split(random_tensor, [1, 2], dim=2) # 返回元组
split_1_2
# (tensor([[[0.0445],
#           [0.2850]],
 
#          [[0.1323],
#           [0.3338]]]),
#  tensor([[[0.0481, 0.1199],
#           [0.1215, 0.0584]],
 
#          [[0.4458, 0.0899],
#           [0.3624, 0.7511]]]))
len(split_1_2)  # 2
split_1_2[0]    # torch.Size([2, 2, 1])
# tensor([[[0.0445],
#          [0.2850]],

#         [[0.1323],
#          [0.3338]]])
split_1_2[1]   # torch.Size([2, 2, 2])
# tensor([[[0.0481, 0.1199],
#          [0.1215, 0.0584]],

#         [[0.4458, 0.0899],
#          [0.3624, 0.7511]]])

所以,torch.split()是一个很有用的函数,可以轻松地将张量分割成任意形状和大小的张量列表,以用于后续处理。

Tips：
感谢@qq_42798074指正
感谢@qq_41720271指正