ESPCN超分复现：亚像素卷积与PixelShuffle原理及PyTorch实现

前几天收拾代码仓库翻到三年前写的ESPCN复现工程。当时为了把论文里的亚像素卷积彻底搞清楚来回折腾了一周连PixelShuffle的输出维度都看反了好几次。现在再看这篇2016年的论文其实是最适合作为超分入门复现的模型之一结构极简三层卷积加一个像素重排就能跑起来但蕴含的思想却影响了后面几乎所有基于CNN的超分模型。这篇博文就把我完整的学习笔记和复现过程整理出来包括原理、代码、踩坑和实测对比希望能帮你省掉一些弯路。1. 为什么2016年的ESPCN至今仍有复现价值1.1 SRCNN给后人留下的效率包袱在ESPCN出现之前超分辨率领域的标杆是SRCNN。它的流程很直白先把低分辨率图像用双三次插值放大到目标尺寸再丢进三层卷积网络里做特征重建。这个思路从效果上看在当时很不错但从工程角度看有一个致命问题——网络从头到尾都在高分辨率空间计算。以×4放大为例如果LR输入是32×32SRCNN的卷积就运行在128×128的特征图上计算量是ESPCN的16倍。而且放大后的图像里大部分像素是插值算法猜出来的冗余信息网络花了大量算力去处理这些本来就不存在的合成像素本质上是在抵消双三次插值留下的模糊痕迹效率很低。ESPCN论文里说得非常直接与其在HR空间做无谓的计算不如先在LR空间把特征提充分最后一步再上采样。这句话后来成了超分模型设计的一条黄金准则。1.2 ESPCN的三个关键设计ESPCN真正让人记住的是三件事的组合LR空间特征提取整个网络的主体部分都在低分辨率空间运行计算复杂度直接降了一个数量级。亚像素卷积层用卷积生成r²×C个通道再通过像素重排periodic shuffling转换为一张r倍分辨率的图像。tanh输出激活最后一层用tanh而不是ReLU输出范围与像素归一化空间[-1,1]对齐。这三个设计是环环相扣的。如果特征提取没有放在LR空间亚像素卷积就失去了低计算量的意义如果最后不用tanh像素重排后的输出可能超出合理像素范围训练时容易震荡。论文里还专门对比了ReLU和tanh结论是在输出层tanh更稳定因为超分本质是回归任务输出理应有界而不是像分类那样需要无界的特征响应。1.3 放到今天还要不要学它这是很多人问我的问题。现在超分领域已经被Transformer、扩散模型刷榜一个2016年的老三层卷积网络还有必要看吗我的答案是非常有必要。亚像素卷积的思想并没有随着ESPCN变老而淘汰它现在依然是很多生成模型默认的上采样方式。不管是GAN的decoder、视频修复模型还是实时的超分应用你几乎都能看到PixelShuffle或它的变体。理解ESPCN相当于拿到了一把理解这些现代模型的钥匙为什么上采样要放到最后为什么r²C通道的重排等价于可学习插值为什么它不会像反卷积那样产生棋盘格伪影。这些底层逻辑搞清楚了看后续论文会顺畅得多。2. 亚像素卷积不是反卷积像素重排的数学直觉2.1 从r²C个通道到高清图我第一次看亚像素卷积的时候被论文里的示意图绕晕了。后来我找到一个特别好的直觉理解方式——反向拆拼图。想象一张HR图像把它按r×r的网格拆开。每个网格位置上的像素可以摊到r²张与LR同尺寸的子图上对应关系是(i, j)位置的相位决定像素落到哪张子图。这个过程是像素重排的逆过程。而亚像素卷积就是把这个逆过程反过来模型在LR空间生成r²张子图然后按相位交错叠起来合成一张r倍大的图像。用形式化语言说输入特征图尺寸是H×W×C最后一层卷积把它映射成H×W×(r²C)PixelShuffle的职责就是把这个张量reshape成(rH)×(rW)×C。通道维度的顺序是关键第k个通道其中k i×r j会对应到HR图像中相位为(i, j)的位置。手动实现一遍能帮助理解def pixel_shuffle_manual(x, r): b, c, h, w x.shape assert c % (r * r) 0 out_c c // (r * r) x x.view(b, out_c, r, r, h, w) # 拆分成 r*r 个相位子图 x x.permute(0, 1, 4, 2, 5, 3).contiguous() # 交错重排 return x.view(b, out_c, h * r, w * r)这段代码跑通了你就真正理解PixelShuffle在做什么了。它不是一个黑盒只是一次精心设计的reshape加permute。2.2 为什么这等价于一种可学习的插值传统插值比如双线性和双三次它们的核心是一个预定义核函数核函数固定不变用LR像素的加权组合造出HR像素。亚像素卷积把同样的逻辑学习化了最后一层3×3卷积同时负责两件事——从LR特征图中提取信息并生成不同相位的子图。每个输出通道的卷积核权重就是对应相位上的插值核。只不过这个核不再是人手设计的固定系数而是从数据里学出来的。这样的好处是不同通道可以学习不同的插值行为。有些通道专门做边缘保持有些通道负责纹理细节恢复有些则处理大块平坦区域的平滑过渡。相比之下固定插值核只有一种行为模式不可能同时兼顾所有这些需求。2.3 反卷积、PixelShuffle与插值的三角关系反卷积转置卷积是另一种可学习上采样方式但它有个臭名昭著的问题当kernel size不能被stride整除时输出会出现棋盘格伪影。原因在于反卷积在数学上是稀疏矩阵运算相邻输出像素接收到的梯度贡献不均匀导致后训练出来的核呈现周期性条纹。亚像素卷积直接从机制上规避了这个问题。它没有做元素填充而是把相位信息分散到了通道维度每一个输出像素都有来自对应相位子图的完整响应激活是平滑的。实际对比生成图像时这个差异非常明显反卷积的结果放大看总有规律的格子纹PixelShuffle的结果则干净很多。PyTorch里nn.PixelShuffle就是官方实现但它要求输入通道数严格等于输出通道数乘以r²。许多人第一次写网络通道数算错运行时报shape错误就是这个原因。3. PyTorch复现的完整代码拆解3.1 网络结构三层卷积加一个PixelShuffle先给出完整网络定义这是我复现后整理出来的最简洁版本import torch import torch.nn as nn class ESPCN(nn.Module): def __init__(self, upscale_factor4): super(ESPCN, self).__init__() self.upscale_factor upscale_factor self.conv1 nn.Conv2d(3, 64, kernel_size5, padding2) self.conv2 nn.Conv2d(64, 64, kernel_size3, padding1) self.conv3 nn.Conv2d(64, 3 * upscale_factor ** 2, kernel_size3, padding1) self.pixel_shuffle nn.PixelShuffle(upscale_factor) self.relu nn.ReLU(inplaceTrue) self.tanh nn.Tanh() def forward(self, x): x self.relu(self.conv1(x)) x self.relu(self.conv2(x)) x self.tanh(self.conv3(x)) x self.pixel_shuffle(x) return x几个细节值得说明conv3的输出通道数为什么是48×4放大时r²16输入是RGB三通道图像16×348。这个数字对应的是为16个相位位置各生成一张三通道子图。第一层用5×5卷积第二第三层用3×3。论文作者试过全3×3效果略差因为第一层直接接触原始LR像素3×3的感受野不足以建模插值所需的局部形状信息。tanh放在conv3之后、PixelShuffle之前。你也可以放在PixelShuffle之后效果几乎一样但按论文的原始设计放前面更合理因为tanh是把每个通道的响应压到[-1,1]像素重排本身只是重新排列不影响值域。3.2 数据准备HR裁剪与LR下采样是复现的关键训练ESPCN通常的做法是从HR图像上随机裁剪固定大小的子块再下采样得到LR。这一步有个工程细节至关重要LR必须由HR经过与测试时完全一致的bicubic下采样得到。很多人直接使用一些库的默认Resize默认插值可能是双线性评估时PSNR会掉一截而且不容易察觉是哪里出了问题。我的数据集逻辑大致是这样from PIL import Image from torch.utils.data import Dataset import torch import numpy as np class SRDataset(Dataset): def __init__(self, hr_paths, scale4, hr_crop_size128): self.hr_paths hr_paths self.scale scale self.hr_crop_size hr_crop_size self.lr_crop_size hr_crop_size // scale def __getitem__(self, idx): hr Image.open(self.hr_paths[idx]).convert(RGB) # 随机裁剪HR子块代码里用numpy做一些索引操作这里略写 hr random_crop(hr, self.hr_crop_size) lr hr.resize((self.lr_crop_size, self.lr_crop_size), Image.BICUBIC) # 归一化到[-1, 1]与tanh输出对齐 lr torch.from_numpy(np.array(lr, dtypenp.float32) / 127.5 - 1.0) hr torch.from_numpy(np.array(hr, dtypenp.float32) / 127.5 - 1.0) return lr.permute(2, 0, 1), hr.permute(2, 0, 1)这里归一化到[-1,1]是刻意的。论文设计tanh输出就意味着网络输出的像素值范围也是[-1,1]训练目标HR必须与它在同一量纲里。如果你习惯用[0,1]那最后一层要么换成sigmoid要么把目标也归一化到[-1,1]千万不要混着来。另外要注意训练时裁剪的HR子块尺寸要能被scale整除否则下采样时出现半个像素的偏移训练集和测试集的分布就错位了。我一般取128×128或192×192对×4来说LR就是32×32或48×48。3.3 训练配置与评估口径训练部分没什么神秘MSE损失加Adam就够了。我的配置如下model ESPCN(upscale_factor4) criterion nn.MSELoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones[200, 400], gamma0.5 ) for epoch in range(500): model.train() for lr, hr in train_loader: optimizer.zero_grad() sr model(lr) loss criterion(sr, hr) loss.backward() optimizer.step() scheduler.step()评估时一定要处理三个口径问题只算Y通道PSNR。RGB转YCbCr后取亮度通道Y计算因为人眼对亮度变化最敏感这是超分任务评估的通行规则。如果直接对RGB三通道算PSNR数值会偏低且不便于和论文对比。去掉边界像素。亚像素卷积在图像边缘处理不完整边界像素会明显拉低指标。常见做法是shave一个scale大小的边框比如×4就裁掉4像素再算。PSNR的max值与归一化方式要对应。归一化到[-1,1]时理论max是2公式是10log10(4/MSE)。归一化到[0,1]时max是1公式是10log10(1/MSE)。两次我都做过用错max值会让数字整体漂移好几dB。下面是我在项目里常用的评估函数def rgb_to_ycbcr(img_tensor): 过程中转Y通道Y 0.257*R 0.504*G 0.098*B 16尺寸需归一化到[0,255]处理 pass # 具体实现省略注意通道顺序即可 def psnr_on_y(sr, hr, shave4): sr_y rgb_to_ycbcr(sr)[:, 0, :, :] hr_y rgb_to_ycbcr(hr)[:, 0, :, :] sr_y sr_y[:, shave:-shave, shave:-shave] hr_y hr_y[:, shave:-shave, shave:-shave] mse ((sr_y - hr_y) ** 2).mean() return 10 * torch.log10(1.0 / mse)这个评估函数里归一化是[0,1]范围所以max值取1.0。如果你的数据跑在[-1,1]记得换成2.0。4. 训练过程中的实测数据与踩坑记录4.1 三个最容易翻车的地方学习率与batch size的匹配。ESPCN结构浅梯度传播路径短Adam配合1e-3在batch size不大比如16的情况下没问题。但如果你为了提高吞吐把batch size开到128学习率必须往下压否则loss曲线会震荡得厉害。我试过1e-3加128的batch前几十个epoch完全像在随机游走降到2e-4才稳住。一个实用经验batch翻倍学习率大概乘0.5。最后一层卷积的初始化与tanh饱和。conv3的输出通道多初始权重方差过大会让tanh直接饱和梯度趋近于零训练前期loss平走好几个epoch。PyTorch默认初始化在多数时候问题不大但如果你用小数据集训练可以在conv3上做一次小方差初始化将权重除以输出通道数的平方根。实践中我观察到默认初始化下前几个epoch的loss下降极慢等激活区域活过来才开始收敛这不算错误但会浪费一点时间。数据增强的对应关系。随机翻转和旋转90度在超分任务里是标配。旋转操作有个隐蔽的坑LR和HR必须做完全相同的变换并且旋转角度只能是90度的整数倍否则插值过程会破坏像素级别的对应关系。如果你用一些库的随机增强比如随机亮度、随机色彩抖动这些对超分训练往往没什么帮助甚至可能让网络学到不期望的颜色偏移建议只用几何变换。4.2 在不同放大倍数下的表现差异我本地分别跑了×2、×4、×8三组实验结论和论文预期一致放大倍数越小ESPCN相对传统插值的优势越明显。×2时三层卷积的感受野足够覆盖局部重建所需范围边缘锐利度提升明显输出几乎看不出伪影。×4时效果依然可观整体比双三次插值高出明显一截放大细节时纹理更自然。×8时问题开始暴露连续放大8倍后需要重建的信息量远超出网络容量平坦区域会出现雾状模糊细纹理直接丢失。这个现象不难解释。×8的等效感受野换算到HR空间只有大约3×(8×3)72像素粗略估计但重建一个8×8区块内的精细纹理需要更大范围的上下文。所以后来超分工作都在往深了堆——SRResNet和EDSR加深度和残差本质就是扩展感受野、增加网络容量。ESPCN的浅层结构注定了它在高倍放大上的天花板。4.3 与SRCNN对比的实验结论我自己完整跑了SRCNN和ESPCN的对比实验严格控制变量相同训练数据、相同epoch数、相同损失函数。结果和论文的定性结论一致。模型特征提取空间上采样方式相对单轮耗时相对测试PSNRSRCNNHR空间预插值卷积约4.2x基线ESPCNLR空间亚像素卷积1.0x约高0.3~0.5dB这张表里我给的是相对量级因为不同实现、不同算力跑出绝对值不一样重点看趋势ESPCN用更短的单轮时间拿到了更高的PSNR。这是ESPCN最有说服力的一点——效率提升没有牺牲质量反而质量略胜。在实际部署场景里这个特性比单纯涨0.5dB更值钱。还有一个体感上的差异ESPCN在×4时的边缘振铃ringing artifact比SRCNN少很多。原因也好理解ESPCN在LR空间做卷积特征图以低频信息为主不会将插值放大后的噪声同步放大输出的HR图像自然更干净。5. 从ESPCN出发亚像素思想在后续工作里的演化5.1 一个核心思想三个方向的演化亚像素卷积发表之后很快成了图像生成领域的基础组件演化大致沿着三个方向展开。第一个方向是上采样模块本身被吸收进各种主干网络。很多生成模型的decoder都用了PixelShuffle或其变体替代反卷积来避免棋盘格伪影。比如图像生成任务里特征图逐步放大时预先设计不同r值的PixelShuffle配合不同通道数已经成为一套标准做法。第二个方向是对亚像素卷积的缺陷做修补。ESPCN原版有一个已知短板相位之间没有显式的交互信息像素重排时子图之间可能不够一致。后来一些工作在这个结构上加注意力机制或可变形卷积让PixelShuffle前的特征图先做自适应的空间重排效果进一步提升。第三个方向是把亚像素思想融合到更宽泛的范式中。比如现在很多连续超分、任意倍率超分模型本质上是用一组基函数或者隐式表示来生成不同相位的特征图最后再用像素重排的方式组装成高分辨率输出。从数学原理看它们和ESPCN是同一个思想家族上采样不是空间放大而是通道到空间的信息重排。5.2 如果重新设计我会怎么改进站在今天的工具链和认知水平上如果让我在ESPCN骨架上做改进我会优先动这几刀加残差连接。ESPCN是纯线性卷积堆叠加一条从输入LR到输出HR的跨尺度残差路径可以先插值再残差相加或者用亚像素卷积后的残差连接能让网络学习细节差而不是完整图像训练会更稳收敛也更快。损失函数换成L1。MSE在超分任务里容易让输出偏平滑因为平方损失对离群大误差的惩罚极重模型宁愿把纹理磨平来降低整体loss。L1损失的梯度更平稳实践中在同架构下普遍能拿到更高的PSNR。在PixelShuffle前加一个轻量注意力模块。亚像素卷积把相位信息分散到通道维度但通道之间没有交互低维度的注意力可以学到一个哪些相位更重要的加权不改变参数量级但对边缘和纹理的保持有帮助。训练策略上先小倍数后大倍数。如果目标是×8直接端到端训练很容易收敛到模糊解。一个实用技巧是分阶段先训练一个×2的子网络再冻结或微调扩展到×4、×8每一级都从合理初始点出发最终效果往往比一次到位好。这些改进单独拎出来都不算有创意但组合在一起其实就是后来很多超分模型的雏形。ESPCN真正留下的财富是它用一种极简的方式证明了上采样可以放进网络这个抽象概念。之后所有关于上采样的设计都是在这个基础上做文章。我在实际复现过程中最深的体会是一个网络的价值不在于它当时刷了多少分而在于它把什么思想留给了后来的人。ESPCN的代码可能只需要一分钟就能敲完但把亚像素卷积的本质吃透之后再去看今天那些复杂的超分模型你会觉得它们都变得透明了。如果你正处于超分入门的阶段强烈建议不要直接跳到最前沿的模型先把这个三层小网络自己从零到一实现一遍这个过程的收获比想象中大得多。