资讯详情

资讯详情

PyTorch实现亚像素卷积:原理、代码与超分实践

把一张低分辨率的图放大到4倍还要让边缘不糊、纹理不崩这事儿远比听起来难。传统做法是先用双三次插值把尺寸撑大再用卷积去“修”可插值引入的模糊根本洗不掉。我最早在某图像处理Demo里就是这么干的放大倍数一高画面就跟蒙了层毛玻璃似的。后来接触到ESPCN那套基于亚像素卷积的方案才意识到问题根源不在后面的修图网络而在上采样这一步本身。这篇文章不绕概念直接用PyTorch把亚像素卷积从原理到代码完整复现一遍涵盖pixel shuffle的数学本质、模型搭建、训练流程以及我实际调参时踩过的坑适合正在做超分辨率、图像生成或任何需要上采样模块的读者。1. 从超分任务说起为什么需要一种“不插值”的上采样方式1.1 传统上采样路子与它们的尴尬做超分辨率的人最早普遍用两种上采样方式。第一种是插值类比如最近邻、双线性和双三次。这类方法本质是根据已知像素推断出未知像素的值可以理解为“猜”。双三次虽然猜得平滑但猜出来的高频信息是假的边缘会有一圈模糊过渡带。关键问题是插值发生在卷积之前后续卷积操作只能在这个已经“糊掉”的大图上做相当于在低质量基础上反复加工天花板很低。第二种是转置卷积也就是常说的反卷积。它把输入像素按stride间隔摆放空出来的位置填零然后用卷积核去扫描。这个方案比插值灵活参数可学习但它有个著名的副作用——棋盘效应。当stride大于1时卷积核的感受野会在输出图上不均匀重叠导致某些区域“多吃”了几次卷积产生明暗条纹。这在高倍放大时特别明显后期只能靠额外的归一化或损失函数压制治标不治本。之所以要单独聊亚像素卷积就是因为它在思路上做了个翻转不先放大图像而是先在低分辨率空间把特征提取充分最后用一次“重排”把多个低分辨率通道合并成一个大图。也就是说上采样不再依赖插值或带重叠的卷积运算而是依赖一个确定性的排列操作。1.2 “亚像素卷积”这个名字其实有误导性很多第一次接触的人会误以为亚像素卷积是一种特殊卷积核。我第一次读相关材料时也这么想后来才明白它真正的核心操作叫周期洗牌英文是periodic shufflingPyTorch里对应torch.nn.PixelShuffle。论文里的核心表达式可以写成[ I^{SR} f^L (I^{LR}) PS(W_L * f^{L-1}(I^{LR}) b_L) ]其中PS就是pixel shuffle操作。整个网络最后几个卷积层把低分辨率特征图变成一个通道数为 (C \times r^2) 的特征图然后PS会把这张图的通道维度重新排列成空间维度直接输出 (r) 倍宽高的大图。换句话说真正的“卷积”部分全部发生在低分辨率空间像素密度低、计算量小放大动作放在最后一个算子并且这个算子不含任何可学习参数只是一个view、permute、view的组合。这也是“亚像素”三个字的由来输入特征图每个像素点对应的其实是输出图一个小区域里的某个亚像素位置网络在学习如何把这些亚像素采样点组合起来。1.3 为什么这种设计在当时是革命性的放在2016年前后的背景下这套设计解决了三个现实问题一是计算量大幅下降因为卷积不作用在高分辨率图上二是避免了转置卷积的棋盘效应因为重排不含重叠运算三是训练更好收敛卷积核可以在低分辨率空间共享特征语义信息提取更充分。超分辨率任务从此有了一个几乎标配的尾部模块。后来SwinIR、Real-ESRGAN这些更强模型虽然主干用Transformer或更深的CNN但在上采样阶段依然选择pixel shuffle足见这个设计的可靠程度。2. 周期性洗牌的原理拆解一次view、一次permute、一次view理解了“过程是重排”还不够你还得知道重排的具体坐标映射否则手写实现时会发现维度总是对不上。这一节我用r2的情况把整个过程拆开讲。2.1 r2时的通道到空间映射假设输入特征图形状是[B, C, H, W]其中C必须是目标通道数C_out的 (r^2) 倍。当r2时输出形状是[B, C_out, 2H, 2W]。从输入到输出每个输出像素都对应一个输入像素的“放大版位置”。具体规则是对于输入张量中属于第c个目标通道的四个子通道索引分别是 (4c0, 4c1, 4c2, 4c3)在输出图中它们分别落在 (2H, 2W)、(2H, 2W1)、(2H1, 2W)、(2H1, 2W1) 这四个相邻位置。用表格表示更直观输入通道索引相对偏移输出位置偏移对应空间位置0(0, 0)左上1(0, 1)右上2(1, 0)左下3(1, 1)右下你可以把输入图看成是输出图的“压缩打包版”四个分辨率减半的通道拼在一起刚好是一张完整的大图。具体拼接顺序有讲究这也是手写实现中最容易出错的地方。2.2 三个步骤的内存语义PyTorch里的PixelShuffle内部逻辑等于下面三步第一步view把通道维度C_out * r * r拆成[C_out, r, r]。这一步本身不改内存顺序只是重新解释形状。第二步permute把维度顺序从[B, C_out, r, r, H, W]调整成[B, C_out, H, r, W, r]。这一步是关键它把“通道子块的行”挪到了“高度维度之后”把“通道子块的列”挪到了“宽度维度之后”。第三步view合并(H, r)得到H*r合并(W, r)得到W*r最终形状变成[B, C_out, H*r, W*r]。为什么第二步permute必须存在因为view只能合并相邻维度。初始形状[B, C_out, r, r, H, W]里r和H不相邻不先通过permute让它们相邻直接view会把不同原始像素的r位置混在一起得到的是乱序图像。2.3 为什么说这是在“亚像素”级别的信息重排传统插值在上采样时创造的是“凭空猜出来的像素”而亚像素卷积重排的是“网络已经学出来的特征”。低分辨率特征图里一个像素点包含的其实是一小块区域的高维表达网络通过卷积层把这块区域的不同细节拆到了不同通道里最后pixel shuffle只是把这个拆解结果按空间位置装回去。可以这么理解普通卷积提取特征时通道维度天然就是“各种频率成分”的容器。亚像素卷积相当于让网络学会把水平偏移、垂直偏移这些信息编码在通道维度上再用一次确定性操作解码回空间维度。解码本身没有参数所以训练时梯度可以通过这个操作直接回传到低分辨率卷积层网络能很自然地学会正确编码。3. 基于PyTorch的实现三种写法与一致性验证3.1 最小可用模型ESPCN风格网络先给一个最简的超分网络定义。输入低分辨率图经过三层卷积提取特征最后用一个卷积把通道数扩到 (C \times r^2)再喂给PixelShuffle。import torch import torch.nn as nn class ESPCN(nn.Module): def __init__(self, upscale_factor4, in_channels3, hidden64): super().__init__() self.upscale_factor upscale_factor self.features nn.Sequential( nn.Conv2d(in_channels, hidden, 5, padding2), nn.Tanh(), nn.Conv2d(hidden, hidden, 3, padding1), nn.Tanh(), nn.Conv2d(hidden, hidden, 3, padding1), nn.Tanh(), ) # 输出通道必须是 in_channels * r^2 self.last_conv nn.Conv2d(hidden, in_channels * upscale_factor ** 2, 3, padding1) self.shuffle nn.PixelShuffle(upscale_factor) def forward(self, x): x self.features(x) x self.last_conv(x) x self.shuffle(x) return x注意last_conv的out_channels计算。目标输出是3通道RGB图r4时最后一个卷积必须输出 (3 \times 16 48) 个通道否则PixelShuffle会报维度错误。3.2 手写pixel_shuffle并与官方算子比对如果你想知道内部到底发生了什么最好手写一遍。下面这个函数和torch.nn.PixelShuffle(r)效果完全一致import torch def manual_pixel_shuffle(x, r): # x shape: [B, C, H, W] B, C, H, W x.shape assert C % (r * r) 0, f通道数 {C} 必须能被 r^2{r*r} 整除 C_out C // (r * r) # 第一步: 把通道维拆成 [C_out, r, r] x x.view(B, C_out, r, r, H, W) # 第二步: 调整顺序让 r 与 H、W 相邻 x x.permute(0, 1, 4, 2, 5, 3).contiguous() # 第三步: 合并相邻的 (H, r) 和 (W, r) x x.view(B, C_out, H * r, W * r) return x验证一致性很简单用随机张量比对即可x torch.randn(2, 48, 8, 8) # r4时输出应为 [2, 3, 32, 32] out1 manual_pixel_shuffle(x, 4) out2 torch.nn.functional.pixel_shuffle(x, 4) print(torch.allclose(out1, out2)) # True很多人在手写时会写成x.view(B, C_out, r, r, H, W)之后直接view不加permute结果输出图是一团乱码。原因就是2.2节说的维度不相邻问题。还有一个高频错误是忘了contiguous()。permute只返回一个共享底层数据的视图维度顺序变了但内存布局没变此时直接view会触发运行时错误必须先contiguous让内存按新顺序重排。3.3 训练循环的基本骨架有了模型训练部分其实和普通图像回归任务没有本质区别。我用一个简化版训练循环说明关键点import torch.optim as optim from torch.utils.data import DataLoader model ESPCN(upscale_factor4).cuda() optimizer optim.Adam(model.parameters(), lr1e-3) criterion nn.L1Loss() for epoch in range(epochs): for lr_imgs, hr_imgs in train_loader: lr_imgs lr_imgs.cuda() hr_imgs hr_imgs.cuda() sr_imgs model(lr_imgs) loss criterion(sr_imgs, hr_imgs) optimizer.zero_grad() loss.backward() optimizer.step()这里一开始就可以用L1Loss而不是很多人习惯的MSE。后文我会具体解释为什么在超分任务里L1通常比L2更好用。4. 完整训练流程数据、损失与评估细节4.1 HR到LR的下采样链路超分训练集的标准构造方式准备一批高分辨率图用双三次插值下采样得到低分辨率图让模型学习从LR恢复HR。这里有个很容易被忽略的细节——下采样前要不要加抗锯齿滤波。torchvision.transforms.Resize里的antialias参数控制这一点。如果直接把HR缩小4倍会存在严重的频率混叠LR里会残留一些伪纹理模型反而能学到不真实的高频先验。更规范的做法是下采样前做一次高斯模糊或使用antialiasTrue保证LR是干净的“理想下采样”结果这样模型学到的才是真正的超分能力而不是隐式去混叠。数据准备代码可以这样写import torchvision.transforms as T # HR - LR 的下采样函数 def hr_to_lr(hr_img, scale4): lr_img T.Resize((hr_img.size[1] // scale, hr_img.size[0] // scale), antialiasTrue)(hr_img) return lr_img训练时建议从HR图里随机裁剪固定大小的patch比如128x128然后下采样得到32x32的LR patch。这样每个batch里模型都在看不同位置的局部结构数据利用率高。再配合随机水平翻转、90度旋转增强能显著抑制过拟合。4.2 损失函数选择L1、L2与Charbonnier超分任务里最经典的损失是MSE因为它直接优化PSNR。但实际训练多了会发现MSE收敛出来的图偏平滑边缘不够锐利。原因是L2损失对离群像素的惩罚是平方级的模型为了避免少量大误差宁愿把边缘糊掉换取整体MSE下降。L1损失的惩罚是线性的对边缘的保留更好。大量实验证明在视觉质量和PSNR之间的平衡上L1往往是最省心的选择。还有一种叫Charbonnier损失本质是L1的光滑近似[ \mathcal{L} \sqrt{(x - y)^2 \epsilon^2} ]其中 (\epsilon) 通常取1e-3。它既保留L1对离群值的鲁棒性又在原点附近可微训练初期不会因为梯度跳变而不稳定。实现也很简单def charbonnier_loss(pred, target, eps1e-3): return torch.mean(torch.sqrt((pred - target) ** 2 eps ** 2))4.3 评估指标PSNR和SSIM的计算坑评估超分模型最常见的是PSNR和SSIM。PSNR的计算很简单但有几个坑像素值范围必须统一。如果模型输出是0到1之间的张量峰值就是1如果是0到255的图峰值就是255。两者算出来的PSNR差很多对比不同论文结果时尤其要小心。PSNR通常在Y通道YCbCr空间的亮度通道上计算而不是RGB三通道直接算。原因是人眼对亮度更敏感且RGB三通道的MSE容易被色度噪声干扰。转换时一般用rgb2ycbcr再取Y通道。SSIM的计算要设置合理的窗口默认11x11高斯窗口在低分辨率图上可能过小建议根据图尺寸调整。评估代码参考import torch.nn.functional as F def compute_psnr(img1, img2, max_val1.0): mse F.mse_loss(img1, img2) return 10 * torch.log10(max_val * max_val / mse)读取测试图时不要用torchvision自带的ToTensor后直接评估要确认精度dtype和值域是否和训练一致。最常见的问题是训练时做了归一化评估时忘了反归一化导致PSNR虚低。5. 对比实验与工程化建议pixel shuffle vs 转置卷积5.1 棋盘效应的根源转置卷积产生棋盘效应根本原因是kernel的重叠不均匀。当stride2、kernel_size3时输出图中某些位置的贡献次数是5次有些位置是3次这种不均匀导致亮度条纹。虽然可以通过初始化卷积核为双线性核来缓解但训练过程中参数一旦偏离初始化效应还是会回来。pixel shuffle不存在这个问题因为重排操作是精确的一对一映射输出图中每个像素都恰好由输入特征图的某一个通道像素贡献没有任何重叠。因此它天生不产生棋盘效应。5.2 计算量的定量对比以一个具体例子说明。假设要把32x32x64的特征图放大到64x64。转置卷积方案卷积核直接作用在填充后的64x64特征图上输入通道64输出通道64一次卷积的乘加次数大约 (64 \times 64 \times 3 \times 3 \times 64 \times 64 \times 64)。pixel shuffle方案先用一个普通卷积把32x32x64变成32x32x(64*r^2)再重排。卷积发生在低分辨率空间空间尺寸只有转置卷积的1/4即使输出通道多了r^2倍总乘加次数仍远小于转置卷积。换算下来pixel shuffle方案在r4时的浮点计算量大约是转置卷积的1/4左右。训练速度差距明显尤其在视频超分这种需要处理海量帧的场景里这个差距会直接决定项目能不能在预算时间内跑完。5.3 什么情况下别用PixelShufflepixel shuffle并非万能。有几种场景我会刻意避开它。一是需要可解释或可控形状的上采样时比如某些语义分割模型的特定解码结构pixel shuffle会破坏空间连续性的先验不如插值后接卷积直观。二是特征图通道数本来就很少时比如中间特征只有8通道r4时最后一个卷积输出的通道数是 (8 \times 16 128)参数爆炸反而不划算。这时候转置卷积或先插值再卷积可能更轻量。三是当上采样倍数不是整数时pixel shuffle没法直接处理。1.5倍、2.5倍这种需求要么先把输入补到合适尺寸要么分段上采样比如1.5倍可以拆成一次3/2倍重排然后裁剪。6. 踩坑记录与调参经验6.1 通道数算错是最常见的错误我在模拟项目X里第一次跑4倍超分网络很快跑起来但训练到一半报错提示pixel shuffle输入通道不匹配。原因是我在last_conv里想当然地写了out_channelshidden忘了乘以upscale_factor ** 2。这个错误之所以隐蔽是因为如果hidden本身是r^2的倍数网络不会立刻报错而是训练十几个epoch后loss异常升高输出图像出现彩色网格。排查方法很简单打印模型最后一层输出的shape确认它是[B, 3 * r * r, H, W]。更稳妥的做法是写一个断言在forward里检查通道数def forward(self, x): x self.features(x) x self.last_conv(x) assert x.size(1) 3 * self.upscale_factor ** 2, \ f通道数 {x.size(1)} 不等于 {3 * self.upscale_factor ** 2} return self.shuffle(x)6.2 BN和PixelShuffle的先后顺序如果你在超分网络里用了BatchNorm要注意BN层和pixel shuffle的先后顺序。pixel shuffle操作不改变batch内每个通道的统计特性但它会改变通道的语义含义重排前通道维度代表的是“不同亚像素位置”重排后通道维度代表的是“不同颜色或特征类型”。如果BN放在pixel shuffle之后BN统计的是重组后特征的分布没问题如果放在之前BN会把同一目标通道的多个子位置作为不同通道独立归一化可能会破坏网络学到的通道间相对关系。我自己的经验是超分任务里很多模型干脆不用BN或者只在特征提取阶段用靠近pixel shuffle的位置尽量不用。原因在于BN对小batch size很敏感而超分训练通常受限于显存batch size往往只有8甚至4BN的统计量不稳定会导致训练震荡。6.3 上采样倍数与通道数的匹配技巧大倍数上采样时通道数的增长是指数级的。5倍超分要求最后卷积输出 (C \times 25) 个通道3通道输出需要75通道3倍超分只要27通道。通道数膨胀带来的不仅是参数增加还有显存峰值上升因为最后一层卷积的输出特征图全部要保存到计算图里用于反向传播。一个实用的技巧是分阶段放大。比如要8倍放大不要直接r8而是两级结构先做2倍pixel shuffle再叠一个2倍pixel shuffle中间加一组卷积做特征整合。这样每级的通道膨胀只有4倍而且中间多了一次特征精修的机会效果往往比一步到位好。6.4 几个值得注意的初始化问题pixel shuffle本身不含参数但前置的那个last_conv的参数初始化会影响收敛速度。PyTorch默认使用Kaiming初始化对ReLU系激活是合理的但如果你像我一样在中间层用了Tanh默认初始化会导致输出方差偏大训练初期梯度爆炸。解决方法是显式初始化def _init_weights(m): if isinstance(m, nn.Conv2d): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) model.apply(_init_weights)还有个小经验最后一层卷积的bias尽量置零。如果bias初始化为随机值pixel shuffle重排时会把这个偏置分散到输出图的各个位置相当于给每块区域叠加了不同直流分量早期图像会出现网格状色偏。6.5 快速验证模型是否正确的“冒烟测试”训练大模型前先用一张小图做冒烟测试。我的做法是构造一个1x3x16x16的随机输入经过模型后检查输出shape是不是1x3x64x64然后跑一次前向和反向确认没有NaN。更有信息量的测试是把输入改成纯色图比如全0此时模型输出应该接近一个常数偏置如果出现明显的棋盘图案基本可以确定pixel shuffle的实现或通道顺序有误。这一类验证放在完整训练之前通常十分钟内就能完成能省下后面排查网络结构错误的好几个小时。写在最后的几点心得做了几次超分相关项目之后我最大的感受是亚像素卷积成功的关键不是那个pixel shuffle算子本身而是它逼迫你把“上采样”这件事放到了最后一步让网络在低分辨率空间专心做特征提取。这个设计思想比它的具体实现更值得借鉴很多上采样场景都可以套用这套“空间不足通道来凑”的思路。如果你在实作中发现最后输出的图片有周期性伪影先别急着换损失函数回头检查一下最后一层卷积输出的通道顺序再检查一下pixel shuffle之后是否有BN层在作怪。这些坑都不难绕过被我列在这一篇里希望能帮你少走几步弯路。另外如果项目跑完还有余力可以试着把pixel shuffle用在扩散模型的上采样路径里你会发现它在生成任务中一样稳。这也算是这套老方法至今没有过时的一个证明。
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →