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