如何在PyTorch中将32×512张量拆分为4个32×128块堆叠为128×128

代码错误原因 #

• torch.split的第一个参数是每个子块的大小，不是拆分的块数：你需要把最后一维（dim=3）的512拆成4个128，所以第一个参数应该填128，你填成4之后相当于拆成了128个长度为4的子块，和需求完全不符
• torch.cat没有指定拼接维度：默认会在第0维拼接，所以直接把batch维度的元素堆了起来，没有在你需要的高度维度（dim=2）拼接

修正后代码 #

import torch 
x = torch.arange(32*512*3*2).reshape(2, 3, 32, 512)
print(x.shape) # 输出 torch.Size([2, 3, 32, 512])

# 按最后一维dim=3拆成4个大小为128的块，每个块形状为 [2,3,32,128]
x_list = x.split(128, dim=3)
print(f"拆分后块数量：{len(x_list)}，每个块形状：{x_list[0].shape}")
# 输出 拆分后块数量：4，每个块形状：torch.Size([2, 3, 32, 128])

# 在dim=2（也就是高度维度）拼接，4个32拼接后得到128
x = torch.cat(x_list, dim=2)
print(f"拼接后形状：{x.shape}")
# 输出 拼接后形状：torch.Size([2, 3, 128, 128])

如果你需要更简洁的实现，也可以不用拆分拼接，直接通过维度变换实现：

x = x.reshape(2, 3, 32, 4, 128).permute(0,1,3,2,4).reshape(2,3,128,128)

内容的提问来源于stack exchange，提问作者DuckQueen