【Torch API】torch.chunk()用法详解

torch.chunk(tensor, chunk_num, dim)与torch.cat()原理相反，它是将tensor按dim（行或列）分割成chunk_num个tensor块，返回的是一个元组。

 a = torch.Tensor([[4,5,7], [3,9,8], [9,6,7]]) b =torch.chunk(a, 3, dim = 1) print(a) print(b)

dim = 0 ,按行分割，行数会变 ；dim = 1 ,按列分割，列数会变  

torch.cat同理，

torch.cat ( (A, B), dim=0)接受一个由两个（或多个）tensor组成的元组，按行拼接，所以两个（多个）tensor的列数要相同。

torch.cat ( (A, B), dim=1)是按列拼接，所以两个tensor的行数要相同。

Bug：IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

代码：

a=torch.tensor([2,1])
b=torch.chunk(a,2,1)
错误原因：

print(a.shape)
torch.Size([2])
a 的 shape 是[2] ，不是[1,2]，是一个一维向量，因此 不存在 dim=1的情况

修改

b=torch.chunk(a,2,0)
b=torch.chunk(a,2,-1)
torch.chunk切割以后还会保留维度