ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch复现SRCNN:三层卷积实现图像超分的关键细节

PyTorch复现SRCNN:三层卷积实现图像超分的关键细节 简介面向PyTorch初学者的SRCNN图像超分完整实现适合刚接触超分辨率或经典卷积网络复现的读者也可作为课程设计、毕业设计的参考项目。资源将数据集、训练好的权重、源码与运行说明整合在同一压缩包内解压后按README操作即可跑通训练或测试流程省去自行准备数据和调通环境的麻烦。压缩包共28个文件、约74.6MB主要包括6个Python脚本数据准备、模型定义、训练与测试、7个pth模型权重、2个h5超分数据集、9张示例bmp图像以及md/txt说明文件结构清晰便于按模块学习。已有2972人学习下载说明该入门方案具备一定参考价值。通过阅读源码与模型文件可掌握SRCNN的搭建思路、训练数据组织方式、Bicubic降采样预处理及超分结果对比方法多个迭代轮次的权重文件和示例效果图也能帮助观察训练收敛过程对于常见参数设置如进程数给出了针对性提醒能降低小白首次运行的门槛。 SRCNN这名字在超分领域算得上“祖师爷”级的存在。我第一次看到它的网络结构时内心是有点失落的三层卷积没有残差没有注意力机制连BatchNorm都没有这不就是最基础的CNN堆叠吗可恰恰是这样一个结构简单到近乎朴素的模型在2014年把图像超分从稀疏编码、字典学习那套手工设计的流程里解放了出来改成了纯粹的端到端回归问题。之后几乎每篇超分论文都会拿它当基线ESPCN、EDSR、RCAN、SRGAN这一长串名字里多多少少都流着SRCNN的血。这篇文章我会用PyTorch把SRCNN整个复现流程走一遍从训练数据怎么造、三层卷积为什么这么设计、到PSNR怎么算才准全部摊开讲适合刚入手PyTorch、第一次接触图像超分的读者也适合想回头补基础的老手。1. SRCNN为什么值得花时间复现一个三层卷积掀起的超分革命1.1 从稀疏编码到端到端学习超分问题的定义说穿了很简单给一张低分辨率图像让你猜它放大后长什么样。难点在于“猜”的这部分低分辨率图像里丢失的高频细节有无数种补法怎么补才能让结果自然、可信这就是算法的核心挑战。在SRCNN之前主流方案是稀疏编码把图像切成小块用预先训练好的字典去表示每个图像块再通过稀疏系数重建高分辨率版本。这流程手工作坊气息很重字典训练、块匹配、系数优化每一步都是一大堆可调参数而且各个环节是分开优化的某个环节调好了整体未必有增益。SRCNN的思路直接得多把整条流程全部丢掉就用三个卷积层逼近“低分辨率图到高分辨率图”这个映射函数。第一层从输入图里提取图像块特征第二层做非线性映射第三层重建像素。三层各司其职对应了传统稀疏编码的三个阶段。端到端训练的好处是所有的中间表示都由数据自己学出来不需要人再手工设计特征规则这也是当年深度学习在很多低层视觉任务上能够碾压传统方法的根本原因。1.2 简单模型反而对理解深度要求更高很多人容易有个误区觉得模型简单就等于复现简单、随便跑跑就能出结果。实际完全不是这样。SRCNN代码确实少核心部分撑死了二三十行但它有四个地方不能想当然数据怎么生成、通道怎么处理、训练输入是什么尺寸、评估时有没有切边界。这四件事每件都不难但任何一个搞错结果就会和论文差出一大截。恰恰因为没有残差、没有注意力那些“挂了就能涨点”的组件兜底整个过程逼着你去把超分的底层逻辑弄清楚。我后来再去看FSRCNN、SRGAN这些模型时经常发现它们的设计动机都是从SRCNN的某个局限出发的。比如FSRCNN针对的是SRCNN输入要先bicubic放大导致计算量大SRGAN针对的是L2损失导致结果过于平滑。没有SRCNN打底这些改进点都无从谈起。这就是为什么我强烈建议每个做超分的人都用PyTorch亲手写一遍SRCNN。2. 训练数据暗藏玄机HR-LR数据对从哪儿来、怎么造2.1 数据不是下载来的是自己“退化”出来的做超分训练的第一件事是搞清楚训练数据长什么样。我们需要的是“低分辨率输入-高分辨率标签”这样的数据对但这世界上没有哪个数据集直接给你配好一一对应的HR-LR对。标准做法是拿高分辨率图像按预设的退化模型降低分辨率得到低分辨率版本。SRCNN当时的退化模型很简单对HR图像做bicubic下采样得到LR然后把LR再bicubic放大回原尺寸作为网络输入网络输出要和最开始的HR一致。这里有个关键点SRCNN是pre-upsampling架构它的输入和输出是同一个尺寸而不是像后来很多模型那样输入小图直接输出大图。所以在数据准备阶段我们必须先把LR放大到HR的尺寸才能喂给网络训练。对应到代码里就是这样import cv2 import numpy as np hr cv2.imread(demo.png, cv2.IMREAD_COLOR) hr cv2.cvtColor(hr, cv2.COLOR_BGR2YCrCb)[:, :, 0] # 保留亮度通道Y h, w hr.shape scale 3 lr cv2.resize(hr, (w // scale, h // scale), interpolationcv2.INTER_CUBIC) lr_up cv2.resize(lr, (w, h), interpolationcv2.INTER_CUBIC) # 关键放大回原尺寸2.2 Dataset和DataLoader的写法和增强细节实际训练时不能一整张图直接丢进网络显存扛不住不说整图的信息冗余也大。通用做法是随机裁剪固定大小的patch比如96x96。裁剪之后再随机做翻转、旋转90度之类的数据增强让一个batch里能看到更多不同的局部结构。我建议直接用cv2统一读写和缩放不要一会儿用PIL一会儿用OpenCV插值算法实现不同会在无形中引入噪声。Dataset这一层可以这么写import torch from torch.utils.data import Dataset import os class SRCNNDataset(Dataset): def __init__(self, hr_dir, scale3, patch_size96): self.hr_paths [os.path.join(hr_dir, f) for f in os.listdir(hr_dir) if f.lower().endswith((.png, .jpg, .bmp))] self.scale scale self.patch_size patch_size def __len__(self): return len(self.hr_paths) def __getitem__(self, idx): hr cv2.imread(self.hr_paths[idx], cv2.IMREAD_COLOR) hr cv2.cvtColor(hr, cv2.COLOR_BGR2YCrCb)[:, :, 0] h, w hr.shape i np.random.randint(0, h - self.patch_size 1) j np.random.randint(0, w - self.patch_size 1) hr_patch hr[i:i self.patch_size, j:j self.patch_size] # 随机水平翻转和90度旋转增加数据多样性 if np.random.rand() 0.5: hr_patch hr_patch[:, ::-1] if np.random.rand() 0.5: hr_patch np.rot90(hr_patch) lr_patch cv2.resize(hr_patch, (self.patch_size // self.scale, self.patch_size // self.scale), interpolationcv2.INTER_CUBIC) lr_up cv2.resize(lr_patch, (self.patch_size, self.patch_size), interpolationcv2.INTER_CUBIC) lr_tensor torch.from_numpy(lr_up.copy()).float().unsqueeze(0) / 255.0 hr_tensor torch.from_numpy(hr_patch.copy()).float().unsqueeze(0) / 255.0 return lr_tensor, hr_tensor2.3 两个容易被忽略的预处理坑第一插值算法必须全局统一。训练时用cv2.INTER_CUBIC生成LR测试时就一定要用同样的设置否则训练和测试面对的退化分布不一致PSNR会凭空掉一截。很多复现不了了之就是死在这种地方。第二归一化范围要想清楚。上面的代码把像素归一化到0~1区间网络输出也就在0~1附近最后画图或算PSNR时得注意还原。我也见过有人训练时不归一化直接拿0~255的值算MSE这样Adam默认的学习率1e-4就显得太小了收敛会很慢。最稳妥的做法就是统一归一化到0~1以后不管换什么模型都沿用这个习惯。3. 模型实现只有十几行代码但每个设计都值得推敲3.1 三层卷积各自的职责SRCNN的经典结构是三段式。第一层是9x9卷积输出64个特征图负责特征提取。9x9的核在当代CNN里算很大了好处是单层就能覆盖足够大的局部区域捕捉图像块的结构信息。第二层是1x1卷积输出32个特征图负责非线性映射本质是把64维特征组合、压缩成更有判别力的32维表示。第三层是5x5卷积输出1个通道负责重建像素。这里的5x5可以理解为在特征空间里做局部聚合最后把高维特征映射回图像像素空间。注意所有卷积都只配ReLU激活没有BatchNorm、没有Dropout这就是SRCNN的全部家底。3.2 PyTorch实现代码非常简单但有两个细节要说明。第一padding的计算必须保证输入输出尺寸一致否则loss都没法算。9x9卷积需要pad41x1不需要5x5需要pad2。第二第一层输入和第三层输出的通道数取决于你在什么色彩空间训练按论文惯例在Y通道上训练就是1进1出直接在RGB上训练就是3进3出。import torch import torch.nn as nn class SRCNN(nn.Module): def __init__(self, in_channels1): super().__init__() self.conv1 nn.Conv2d(in_channels, 64, kernel_size9, padding4) self.conv2 nn.Conv2d(64, 32, kernel_size1, padding0) self.conv3 nn.Conv2d(32, in_channels, kernel_size5, padding2) self.relu nn.ReLU(inplaceTrue) def forward(self, x): x self.relu(self.conv1(x)) x self.relu(self.conv2(x)) x self.conv3(x) return x3.3 为什么训练输入是“放大后的LR”而不是原始LR这个问题我问过不少人答不上来的占多数。SRCNN不属于直接在低分辨率空间做计算的模型它先把LR用bicubic插值放大到目标分辨率再让卷积网络去学习和HR之间的残差关系。这个设计今天看有点“浪费算力”因为放大后的图尺寸很大特征图也大网络却要在上面做全部卷积计算。但放在2014年这个选择很自然先把大部分工作交给传统插值做网络只负责“修补”放大后依然模糊的细节问题被简化成了输入输出尺寸一致的回归问题模型设计和训练都轻松很多。后来FSRCNN、ESPCN改用后置上采样把特征提取先放在LR空间完成效率更高但那是后话。复现SRCNN时不要自作聪明把LR直接塞进去然后期望输出HR尺寸那是另一种模型的玩法硬套在这里只会让尺寸对不上代码直接报错。4. 训练流程没多复杂但损失函数、学习率和评估方式里面全是细节4.1 损失函数为什么选L2而不是L1SRCNN的训练目标是让重建图与HR图之间的MSE最小化对应的损失就是L2损失criterion nn.MSELoss()为什么不选L1因为在超分这个传统任务里评价指标PSNR和MSE直接相关PSNR 10 * log10(MAX^2 / MSE)。你优化MSE就是在直接优化PSNR。虽然L2损失会让输出偏向平滑、缺少锐利细节后来的SRGAN就是拿这个说事但在SRCNN那个年代L2就是最合理的选择。训练用的数据是YCbCr色彩空间下的Y通道也就是灰度图所以是个单通道回归问题难度更低收敛也快。4.2 优化器、学习率与训练循环原论文用的是SGD加动量学习率1e-4。实际复现中用Adam替换也没问题收敛往往更快更省心。SRCNN参数量很小只有几万级别显存占用很友好batch size设64甚至更大都很轻松。我习惯把学习率设置为1e-4每训练30个epoch衰减一半配合StepLR调度器。epoch数不用设很大我用DIV2K里抽出的几十张图训练大概100个epoch就能看到PSNR明显上扬。训练循环是非常标准的import torch.optim as optim from torch.utils.data import DataLoader from torch.optim.lr_scheduler import StepLR model SRCNN().cuda() optimizer optim.Adam(model.parameters(), lr1e-4) scheduler StepLR(optimizer, step_size30, gamma0.5) criterion nn.MSELoss() dataset SRCNNDataset(data/hr, scale3, patch_size96) loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers4) for epoch in range(100): model.train() total_loss 0.0 for lr, hr in loader: lr, hr lr.cuda(), hr.cuda() sr model(lr) loss criterion(sr, hr) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() scheduler.step() print(fEpoch {epoch 1}, Loss: {total_loss / len(loader):.6f})4.3 验证PSNR时最容易翻车的细节每个epoch结束我都会在Set5或Set11这种标准测试集上算一次PSNR用来看训练是否真的在变好。注意这里有个经典坑用卷积做重建图像边界区域会因为padding产生误差直接计算会把PSNR拉低一些。论文里评估时通常会忽略边界像素一般做法是左右上下各切掉6个像素再算这个“shave border”的细节直接决定了你的指标能不能对得上论文。完整的PSNR计算函数我这样写的def psnr(img1, img2, shave_border6): if shave_border 0: img1 img1[shave_border:-shave_border, shave_border:-shave_border] img2 img2[shave_border:-shave_border, shave_border:-shave_border] mse np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2) if mse 0: return float(inf) max_pixel 255.0 if img1.max() 1.0 else 1.0 return 10 * np.log10(max_pixel ** 2 / mse)你可能会遇到一个很诡异的局面训练loss在降PSNR却不涨。十有八九是评估代码里输入图像的范围和标签不在同一个尺度下比如网络输出是0~1的浮点图却拿它和0~255的标签去比结果一片混乱。我见过身边不少人踩这个坑排查半天最后发现就是把结果乘回255就能解决。5. 测试推理与结果核对不切线算PSNR等于白算5.1 完整的推理链路Y通道重建CbCr通道放大就好超分推理和训练要严格保持同一个流程。RGB图像先转成YCbCr只对Y通道做SRCNN重建CbCr两个色度通道直接bicubic放大到目标尺寸最后再把三个通道合并转回RGB。原因很朴素人眼对亮度信息最敏感超分指标PSNR也是在Y通道上算的色度通道本来就信息量低不值得网络去专门建模。如果你直接对RGB三个通道做SRCNN问题就是模型要同时学三个通道的分布训练难度变大而且最终评估时还得转成YCbCr再算Y通道PSNR兜了一圈反而不便。按论文惯例走Y通道省时省力对比结果也方便。推理部分可以这样写注意一定用torch.no_grad()避免显存和计算开销def infer(model, lr_img, scale3): lr_ycc lr_img.convert(YCbCr) y, cb, cr lr_ycc.split() y_up y.resize((lr_img.width * scale, lr_img.height * scale), Image.BICUBIC) y_tensor torch.from_numpy(np.array(y_up, dtypenp.float32)).unsqueeze(0).unsqueeze(0) / 255.0 model.eval() with torch.no_grad(): sr_y model(y_tensor.cuda()) sr_y sr_y.squeeze(0).squeeze(0).cpu().clamp(0, 1) * 255.0 sr_y sr_y.numpy().astype(np.uint8) sr_y Image.fromarray(sr_y, modeL) cb_up cb.resize((lr_img.width * scale, lr_img.height * scale), Image.BICUBIC) cr_up cr.resize((lr_img.width * scale, lr_img.height * scale), Image.BICUBIC) sr_img Image.merge(YCbCr, (sr_y, cb_up, cr_up)).convert(RGB) return sr_img5.2 复现到什么程度算成功参考指标对照很多新手复现完会盯着论文的PSNR数字发呆生怕差0.1个dB就是自己写错了。我的经验是差的只要控制在合理范围内不用太纠结。下面是论文公开结果和一般复现能达到的区间训练数据沿用常见公开数据时大致如此测试集放大倍数论文参考PSNRY通道一般复现区间Set5×236.66 dB36.3 ~ 36.7 dBSet5×332.75 dB32.4 ~ 32.8 dBSet5×430.49 dB30.1 ~ 30.5 dBSet14×330.23 dB29.9 ~ 30.3 dB如果你的结果比论文低了0.5dB以上建议先从四个方向排查训练集的bicubic是否和测试集一致、评估有没有切边界、归一化是否还原、训练数据量是否足够。如果只是低了0.2dB八成是训练数据选择不同导致的不用太焦虑。5.3 最容易踩的四个问题我把自己和身边人踩过的坑整理成一份速查清单插值实现不一致训练和推理必须用同一个插值库建议统一cv2。归一化混用0~1和0~255的图混着比较或者使用0~255的数据配1e-4学习率都会导致指标异常。统一0~1就行。通道没对齐只用Y通道训练时测试环节不要忘了把CbCr放大后合并回去。很多人只超分了Y然后直接把单通道图保存为灰度图看起来像那么回事实际上丢了色度信息。训练集太少SRCNN虽然简单但也不是几张图就能学明白的。至少准备几十张内容丰富的HR图裁剪patch数量不够的话很容易在验证集上过拟合。最后再说一点个人体会。我第一次做这个复现时偷懒直接拿RGB三通道训练测Set5比论文低快0.5dB还以为是网络实现有问题查了很久才发现问题出在评估口径不同。论文是按Y通道报的PSNR我在RGB空间算平均数字当然对不上。后来老老实实切到YCbCr流程改用Y通道训练和测试指标一下就回到了正常范围。这个流程说来简单但很多人复现时不看论文implementation details就容易在这种地方卡住。把SRCNN完整跑通之后建议你做两个小实验巩固理解一是把9x9卷积换成3x3堆叠两层看看PSNR怎么变二是把中间层1x1卷积去掉让64维特征直接进5x5重建层看看又是什么效果。这两个小改动都能帮你更快建立对超分模型设计的直觉。等这层纸捅破再看后面的超分模型你会觉得它们的思路都不难理解。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进