☰
PyTorch实现亚像素卷积:原理、代码与超分实践
2026/10/18 23:56:52 网站建设 项目流程

把一张低分辨率的图放大到4倍,还要让边缘不糊、纹理不崩,这事儿远比听起来难。传统做法是先用双三次插值把尺寸撑大,再用卷积去“修”,可插值引入的模糊根本洗不掉。我最早在某图像处理Demo里就是这么干的,放大倍数一高,画面就跟蒙了层毛玻璃似的。后来接触到ESPCN那套基于亚像素卷积的方案,才意识到问题根源不在后面的修图网络,而在上采样这一步本身。这篇文章不绕概念,直接用PyTorch把亚像素卷积从原理到代码完整复现一遍,涵盖pixel shuffle的数学本质、模型搭建、训练流程,以及我实际调参时踩过的坑,适合正在做超分辨率、图像生成或任何需要上采样模块的读者。

1. 从超分任务说起:为什么需要一种“不插值”的上采样方式

1.1 传统上采样路子与它们的尴尬

做超分辨率的人,最早普遍用两种上采样方式。

第一种是插值类,比如最近邻、双线性和双三次。这类方法本质是根据已知像素推断出未知像素的值,可以理解为“猜”。双三次虽然猜得平滑,但猜出来的高频信息是假的,边缘会有一圈模糊过渡带。关键问题是,插值发生在卷积之前,后续卷积操作只能在这个已经“糊掉”的大图上做,相当于在低质量基础上反复加工,天花板很低。

第二种是转置卷积,也就是常说的反卷积。它把输入像素按stride间隔摆放,空出来的位置填零,然后用卷积核去扫描。这个方案比插值灵活,参数可学习,但它有个著名的副作用——棋盘效应。当stride大于1时,卷积核的感受野会在输出图上不均匀重叠,导致某些区域“多吃”了几次卷积,产生明暗条纹。这在高倍放大时特别明显,后期只能靠额外的归一化或损失函数压制,治标不治本。

之所以要单独聊亚像素卷积,就是因为它在思路上做了个翻转:不先放大图像,而是先在低分辨率空间把特征提取充分,最后用一次“重排”把多个低分辨率通道合并成一个大图。也就是说,上采样不再依赖插值或带重叠的卷积运算,而是依赖一个确定性的排列操作。

1.2 “亚像素卷积”这个名字其实有误导性

很多第一次接触的人会误以为亚像素卷积是一种特殊卷积核。我第一次读相关材料时也这么想,后来才明白:它真正的核心操作叫周期洗牌,英文是periodic shuffling,PyTorch里对应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

理解了“过程是重排”还不够,你还得知道重排的具体坐标映射,否则手写实现时会发现维度总是对不上。这一节我用r=2的情况把整个过程拆开讲。

2.1 r=2时的通道到空间映射

假设输入特征图形状是[B, C, H, W],其中C必须是目标通道数C_out的 (r^2) 倍。当r=2时,输出形状是[B, C_out, 2H, 2W]。

从输入到输出,每个输出像素都对应一个输入像素的“放大版位置”。具体规则是:对于输入张量中属于第c个目标通道的四个子通道(索引分别是 (4c+0, 4c+1, 4c+2, 4c+3)),在输出图中它们分别落在 (2H, 2W)、(2H, 2W+1)、(2H+1, 2W)、(2H+1, 2W+1) 这四个相邻位置。用表格表示更直观:

输入通道索引(相对偏移)输出位置偏移对应空间位置
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_factor=4, in_channels=3, hidden=64): super().__init__() self.upscale_factor = upscale_factor self.features = nn.Sequential( nn.Conv2d(in_channels, hidden, 5, padding=2), nn.Tanh(), nn.Conv2d(hidden, hidden, 3, padding=1), nn.Tanh(), nn.Conv2d(hidden, hidden, 3, padding=1), nn.Tanh(), ) # 输出通道必须是 in_channels * r^2 self.last_conv = nn.Conv2d(hidden, in_channels * upscale_factor ** 2, 3, padding=1) 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图,r=4时,最后一个卷积必须输出 (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) # r=4时,输出应为 [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_factor=4).cuda() optimizer = optim.Adam(model.parameters(), lr=1e-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里会残留一些伪纹理,模型反而能学到不真实的高频先验。更规范的做法是下采样前做一次高斯模糊或使用antialias=True,保证LR是干净的“理想下采样”结果,这样模型学到的才是真正的超分能力,而不是隐式去混叠。

数据准备代码可以这样写:

import torchvision.transforms as T # HR -> LR 的下采样函数 def hr_to_lr(hr_img, scale=4): lr_img = T.Resize((hr_img.size[1] // scale, hr_img.size[0] // scale), antialias=True)(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, eps=1e-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_val=1.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的重叠不均匀。当stride=2、kernel_size=3时,输出图中某些位置的贡献次数是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方案在r=4时的浮点计算量大约是转置卷积的1/4左右。训练速度差距明显,尤其在视频超分这种需要处理海量帧的场景里,这个差距会直接决定项目能不能在预算时间内跑完。

5.3 什么情况下别用PixelShuffle

pixel shuffle并非万能。有几种场景我会刻意避开它。

一是需要可解释或可控形状的上采样时,比如某些语义分割模型的特定解码结构,pixel shuffle会破坏空间连续性的先验,不如插值后接卷积直观。

二是特征图通道数本来就很少时,比如中间特征只有8通道,r=4时最后一个卷积输出的通道数是 (8 \times 16 = 128),参数爆炸,反而不划算。这时候转置卷积或先插值再卷积可能更轻量。

三是当上采样倍数不是整数时,pixel shuffle没法直接处理。1.5倍、2.5倍这种需求,要么先把输入补到合适尺寸,要么分段上采样,比如1.5倍可以拆成一次3/2倍重排然后裁剪。

6. 踩坑记录与调参经验

6.1 通道数算错是最常见的错误

我在模拟项目X里第一次跑4倍超分,网络很快跑起来,但训练到一半报错,提示pixel shuffle输入通道不匹配。原因是我在last_conv里想当然地写了out_channels=hidden,忘了乘以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甚至4,BN的统计量不稳定,会导致训练震荡。

6.3 上采样倍数与通道数的匹配技巧

大倍数上采样时,通道数的增长是指数级的。5倍超分要求最后卷积输出 (C \times 25) 个通道,3通道输出需要75通道,3倍超分只要27通道。通道数膨胀带来的不仅是参数增加,还有显存峰值上升,因为最后一层卷积的输出特征图全部要保存到计算图里用于反向传播。

一个实用的技巧是分阶段放大。比如要8倍放大,不要直接r=8,而是两级结构:先做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用在扩散模型的上采样路径里,你会发现它在生成任务中一样稳。这也算是这套老方法至今没有过时的一个证明。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询