torch.chunk(tensorchunk_numdim)与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切割以后还会保留维度

 

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐