Pytorch中的split()和chunk()函数功能和使用区别

引 言 torch.split()函数和torch.chunk()都是用来对数据按照某个维度进行切分，语法形式类似但是参数实际表达含义略有不同。本文通过具体使用案例生动说明两个函数的区别。

---

一、torch.split()

语法：

split( tensor: Tensor, split_size_or_sections: Union[int, List[int]], dim: int = 0
) -> List[Tensor]:

• tensor:输入数据张量
• split_size_or_sections：可以为整数代表以个张量块的切分大小，若张量不够切分，把剩余部分当作一块处理。也可以为列表，代表每块大小，列表和一定于数据量相同。
• dim:确定切分维度
• List[Tensor]：输出切分后的张量列表

示例：

1. 在维度1上切分，每块大小为2。

data = torch.randn(2,6,10)
result_split = torch.split(data,2,dim=1) 	#也可写成 result_split = data.split(2,dim=1)
for e in result_split:
    print(e.shape)
   
###显示结果###
torch.Size([2, 2, 10])
torch.Size([2, 2, 10])
torch.Size([2, 2, 10])

2. 在维度1上切分，每块大小为4，最后剩余部分组成一块。

data = torch.randn(2,6,10)
result_split = data.split(4,dim=1)

for e in result_split:
    print(e.shape)
    
###显示结果###
torch.Size([2, 4, 10])
torch.Size([2, 2, 10])

3. 在维度1上切分，切分大小为一个列表，每块按照列表中的大小确定。注：列表元素和一定与该维度数据相同。

data = torch.randn(2,6,10)
result_split = data.split([2,1,3],dim=1)
for e in result_split:
    print(e.shape)

###显示结果###
torch.Size([2, 2, 10])
torch.Size([2, 1, 10])
torch.Size([2, 3, 10])

二、torch.chunk()

语法：

chunk(input: Tensor, chunks: _int, dim: _int=0) -> List[Tensor]:

• input：输入数据张量
• chunks：只能为整数，指定数据切分多少块
• dim:确定切分维度
• List[Tensor]：输出切分后的张量列表

示例：
在维度1上切分，总共切分为2块。

data = torch.randn(2,6,10)
result_chunk = data.chunk(2,dim=1)
for e in result_chunk:
    print(e.shape)
    
###显示结果###
torch.Size([2, 3, 10])
torch.Size([2, 3, 10])

三、torch.split()和torch.chunk()的区别

1. torch.split()中的split_size_or_sections代表每块大小。torch.chunk()中的chunks代表切分成多少块。
2. 当切分维度上的数据小于切分大小或切分块数时，torch.chunk()在该维度切分成块大小为1的数据而torch.split()则不会切分。

示例：

# torch.chunk()
data = torch.randn(2,6,10)
result_chunk = data.chunk(7,dim=1)
for e in result_chunk:
    print(e.shape)

###显示结果###
torch.Size([2, 1, 10])
torch.Size([2, 1, 10])
torch.Size([2, 1, 10])
torch.Size([2, 1, 10])
torch.Size([2, 1, 10])
torch.Size([2, 1, 10])

# torch.split()
data = torch.randn(2,6,10)
result_split = data.split(7,dim=1)
for e in result_split:
    print(e.shape)

###显示结果###
torch.Size([2, 6, 10])