将批量单通道图像展开为指定尺寸Patch的PyTorch实现问题

你当前代码的问题主要有三点：

1. 输入张量维度不符合Unfold要求 #

torch.nn.Unfold的输入必须是4维张量(N, C, H, W)，其中C是通道数。你的原始张量是(4,180,320)，缺少通道维度，Unfold会错误地把180当成通道数C，导致输出的第一个维度变成180×64×64=16384，这和你预期的完全不符。

2. 未处理0填充 #

你的图像高度180无法被64整除（180=64×2+52），不填充的话只能取2行×5列=10个patch，但你需要先把图像填充到能被patch_size整除的尺寸（比如256×320，64×4=256），才能得到完整的所有patch。

3. Unfold输出格式需要重新调整 #

Unfold的输出形状是(N, C×h_p×w_p, p)，你需要转置并reshape才能得到目标的(N, p, h_p, w_p)格式。

修正后的代码实现 #

import torch
import torch.nn as nn

# 原始张量：(N, H, W) → 补充通道维度转为(N, C, H, W)
tensor = torch.randn(4, 180, 320).unsqueeze(1)  # 现在形状为(4, 1, 180, 320)
patch_size = (64, 64)
h_p, w_p = patch_size

# 计算需要填充的像素数：补到最近的能被patch_size整除的尺寸
h, w = tensor.shape[2], tensor.shape[3]
pad_h = (h_p - h % h_p) % h_p  # 若h能整除则pad_h=0，否则补差值
pad_w = (w_p - w % w_p) % w_p

# 执行0填充：ZeroPad2d参数顺序为(left, right, top, bottom)
pad_layer = nn.ZeroPad2d((0, pad_w, 0, pad_h))
padded_tensor = pad_layer(tensor)  # 填充后形状为(4, 1, 256, 320)

# 执行Unfold操作
unfold = nn.Unfold(kernel_size=patch_size, stride=patch_size)
unfolded = unfold(padded_tensor)  # 形状为(4, 64*64, 20)，其中20=4行×5列

# 调整为目标形状：(N, p, h_p, w_p)
target_tensor = unfolded.transpose(1, 2).reshape(4, -1, h_p, w_p)
print(target_tensor.shape)  # 输出 torch.Size([4, 20, 64, 64])

关于你自行实现的方案检查建议 #

如果你的方案存在以下情况，可能有错误或性能问题：

• 未补充通道维度直接使用Unfold：会导致维度计算混乱，patch内容错误
• 填充计算错误：比如只填充单侧或计算的填充量不对，导致patch边缘内容异常
• 直接reshape而非先转置：Unfold输出的第二个维度是展平的patch，直接reshape会打乱patch的空间顺序
• 手动循环分块：性能远低于Unfold的底层优化实现，尤其是批量处理时

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