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切割以后还会保留维度