简介SRCNN 是图像超分辨率领域的经典卷积网络这个 PyTorch 实现面向刚接触超分辨率方向的初学者与课程实践者把数据准备、模型定义、训练与测试的完整流程跑通用来解决论文复现时环境配置繁琐、数据格式不统一的问题。压缩包共 28 个文件、约 74.6MB其中 6 个 py 脚本分别负责数据预处理、网络搭建、训练与推理7 个 pth 权重文件包含不同 epoch 的模型和最优模型另外还有 h5 格式的训练集与测试集数据、若干 bmp 结果对比图以及 README 说明文档解压后按目录即可对照查阅。已有 2978 人学习下载属于上手门槛较低的入门级复现项目。读者能直接拿到可运行的训练与推理代码、现成权重与数据集观察双三次插值与 SRCNN 放大结果在蝴蝶、斑马、文字等样本上的清晰度差异并借助 README 中的参数注释理解进程数、路径等常见报错的排查思路为后续更换放大倍率或自建数据继续实验打下基础。1. 一张低清图放大后全是锯齿SRCNN 到底在补什么把一张 320×240 的截图放大三倍双三次插值能交差但文字边缘和纹理发虚。图像超分要补的是高频细节不是简单拉伸像素。SRCNN 是超分辨率里最早把三层卷积端到端接起来的模型之一先双三次插值到目标尺寸再在 YCbCr 的 Y 通道上做 9-5-5 卷积映射。它参数少、结构清楚适合想跑通第一套 PyTorch 超分训练管线的开发者也适合拿它当 baseline 验证数据管线和评估代码。别指望它今天能打 SOTA但把它跑顺后面的 SR 模型拆起来会轻松很多。2. 项目骨架与数据管线从 HR 图到 YCbCr 单通道输入的完整准备2.1 目录结构与依赖安装这套实现按“数据、模型、工具、入口”四块拆目录不花哨但每个文件都有明确职责。训练入口只负责解析参数和调度模型文件只放网络结构数据文件只处理 HR 到 LR 的转换推理脚本单独放避免训练和推理互相 import 出一堆循环依赖。下面是我习惯的目录形态源码包拿到手可以先对照一遍srcnn_project/ ├── data/ │ ├── HR/ # 高清原图训练和验证都从这里切 │ └── val/ # 可选单独放验证图按文件名划分 ├── models/ │ └── srcnn.py # SRCNN 网络定义与权重初始化 ├── utils/ │ ├── dataset.py # HR 切块、下采样、YCbCr 转换 │ ├── metrics.py # PSNR 计算 │ └── image.py # 推理切块、拼接、保存 ├── train.py # 训练入口 ├── inference.py # 单图/批量推理入口 ├── requirements.txt └── checkpoints/ # 训练产生的 best.pth、last.pthmodels/srcnn.py只放网络不碰数据utils/dataset.py里不要写训练循环train.py里不要塞网络层定义。这样后面把 SRCNN 换成别的结构时数据管线和评估代码不用动。依赖也不多核心就是 PyTorch、Pillow、NumPy 和 tqdm安装时按本机 CUDA 选对应的 torch 轮子别在 CPU 上硬扛大图训练。pip install torch torchvision pillow numpy tqdmtorchvision在这里主要用ToTensor也可以手写img / 255.0后转 tensor少一个依赖。tqdm只负责进度条不影响训练逻辑。第一次跑建议先把batch_size调到 4、patch_size调到 64用十几张图确认前向、反向、保存权重都通再上全量数据。血泪经验是数据管线没验证就直接开 100 epoch最后 loss 不降回头查半天才发现是通道顺序或下采样顺序反了。2.2 训练集切块、下采样与色彩空间转换SRCNN 的经典做法是在 YCbCr 的 Y 通道上训练CbCr 通道用双三次插值放大后合并。原因不玄学人眼对亮度细节更敏感Y 通道集中了大部分结构信息只训一个通道也能把参数量和显存压下来。数据管线要干三件事从 HR 图随机裁 patch、把 HR 的 Y 通道下采样成 LR、再把 LR 双三次上采样回 HR 尺寸作为网络输入。注意顺序不能反网络输入是“已经放大但模糊”的 Y标签是 HR 的 Y。import random from pathlib import Path from PIL import Image from torch.utils.data import Dataset from torchvision.transforms.functional import to_tensor class SRDataset(Dataset): def __init__(self, hr_dir, patch_size96, scale3, trainTrue): self.hr_paths sorted([p for p in Path(hr_dir).rglob(*) if p.suffix.lower() in {.png, .jpg, .jpeg}]) self.patch_size patch_size self.scale scale self.train train def __len__(self): return len(self.hr_paths) def __getitem__(self, idx): hr Image.open(self.hr_paths[idx]).convert(RGB) w, h hr.size if self.train: # 随机裁 HR patch避免网络只记住固定位置 if w self.patch_size or h self.patch_size: hr hr.resize((max(w, self.patch_size), max(h, self.patch_size)), Image.BICUBIC) w, h hr.size x random.randint(0, w - self.patch_size) y random.randint(0, h - self.patch_size) hr hr.crop((x, y, x self.patch_size, y self.patch_size)) # 水平翻转增强超分任务里别做垂直翻转过度增强 if random.random() 0.5: hr hr.transpose(Image.FLIP_LEFT_RIGHT) else: # 验证集用中心裁剪保证每次评估输入一致 left max(0, (w - self.patch_size) // 2) top max(0, (h - self.patch_size) // 2) hr hr.crop((left, top, left self.patch_size, top self.patch_size)) hr_ycbcr hr.convert(YCbCr) y, cb, cr hr_ycbcr.split() # HR Y - LR Y - 双三次上采样回 HR 尺寸作为网络输入 lr_w, lr_h y.size[0] // self.scale, y.size[1] // self.scale lr_y y.resize((lr_w, lr_h), Image.BICUBIC) lr_y_up lr_y.resize(y.size, Image.BICUBIC) return to_tensor(lr_y_up), to_tensor(y)patch_size是 HR 侧裁出来的边长scale是放大倍数。scale3时 HR patch 96×96LR 是 32×32再上采样回 96×96。trainTrue时做随机裁剪和水平翻转验证时用中心裁剪保证 PSNR 可比。to_tensor会把 0-255 转成 0-1 的 float后面损失函数和 PSNR 都按这个范围算。如果源码包里 Dataset 返回的是 RGB 三通道先别急着改网络先把convert(YCbCr)和split()补上否则训练也能跑但指标和视觉结果都对不上经典 SRCNN 的预期。提示验证集不要做随机裁剪否则每轮 PSNR 抖动很大你很难判断模型到底有没有变好。2.3 一次前向传播的张量形状推导把形状写清楚后面改网络时不容易翻车。以scale3、patch_size96为例DataLoader 一个 batch 出来是[B, 1, 96, 96]标签也是[B, 1, 96, 96]。SRCNN 三层卷积都带 padding输出空间尺寸保持不变阶段输入形状卷积配置输出形状说明输入[B,1,96,96]无[B,1,96,96]双三次上采样后的 YConv1[B,1,96,96]9×9, 64, padding4[B,64,96,96]提取局部纹理Conv2[B,64,96,96]5×5, 32, padding2[B,32,96,96]非线性映射Conv3[B,32,96,96]5×5, 1, padding2[B,1,96,96]重建 Y参数量大约 5.7 万第一层9×9×1×6464第二层5×5×64×3232第三层5×5×32×11。感受野是9 (5-1) (5-1)17所以 HR patch 至少得比 17 大96 是很保守的尺寸。patch 太小边界像素占比高训练会学到一堆补零边缘patch 太大显存吃得快而且 SRCNN 本身容量小收益不明显。我一般从 96 开始显存够就上 128不够就降到 64但别低于 32。3. 模型实现三层卷积 9-5-5 的 SRCNN 网络拆解3.1 网络定义与权重初始化SRCNN 的结构没什么花活三层卷积加两层 ReLU最后一层不加激活直接回归 Y 值。第一层 9×9 大核抓低频结构第二层 5×5 做非线性映射第三层 5×5 重建。padding分别给 4、2、2保证输出和输入同尺寸。下面这份定义可以直接放进models/srcnn.pyimport torch import torch.nn as nn class SRCNN(nn.Module): def __init__(self): super().__init__() self.net nn.Sequential( nn.Conv2d(1, 64, kernel_size9, padding4), nn.ReLU(inplaceTrue), nn.Conv2d(64, 32, kernel_size5, padding2), nn.ReLU(inplaceTrue), nn.Conv2d(32, 1, kernel_size5, padding2) ) def forward(self, x): # x: [B, 1, H, W]已经是双三次上采样后的 Y return self.net(x) def init_weights(m): if isinstance(m, nn.Conv2d): # ReLU 前用 Kaiming normal比全零初始化稳 nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: nn.init.zeros_(m.bias)inplaceTrue的 ReLU 省一点显存但在需要保留中间特征做可视化时别用免得把前一层输出改掉。权重初始化用 Kaiming normalfan_out对 ReLU 比较合适偏置置零。原始 SRCNN 论文里用的是更小方差的高斯初始化放到 PyTorch 里不必照抄Kaiming 已经够稳。如果训练几轮 loss 就 NaN先查输入有没有归一化到 0-1再查学习率是不是太大最后才怀疑初始化。3.2 上采样位置与感受野为什么先插值再卷积SRCNN 把双三次上采样放在网络外面模型只处理 HR 尺寸的 Y。这和后来很多“在模型里做上采样”的做法不一样。好处是实现简单输入输出同尺寸损失函数直接逐像素比坏处是卷积在 HR 尺寸上算大图推理显存吃得多。读代码时看到forward里没有Upsample或PixelShuffle是正常的上采样在Dataset和推理脚本里完成。这个选择对二次开发有影响如果你想改成残差学习让网络只学高频输入仍然是双三次上采样后的 Y输出x residual改动很小。如果你想换成 ESPCN 那种“低分辨率特征 亚像素卷积”数据管线就得把 LR Y 直接喂进来不能再预先上采样。两条路别混着走混着走最常见的翻车是数据管线已经上采样一次模型里又上采样一次输出尺寸翻倍loss 直接 shape mismatch。感受野 17 也决定了 patch 不能太小。推理整图时边缘像素的卷积会补零拼块时如果直接硬拼接缝处会有轻微亮暗差。后面第 4 章会用 overlap 切块取中心区域把这个边界效应压下去。训练阶段用随机裁剪其实也在让网络见过各种边界位置算是一种廉价的数据增强。3.3 损失函数、优化器与训练循环损失用 MSE优化器用 Adam学习率从1e-4起步。SRCNN 参数量小Adam 比 SGD 更容易在短周期内看到效果如果你要严格复现论文可以换 SGD 动量但调参成本更高。训练循环里每个 batch 只做前向、算损失、反向、更新验证阶段算 PSNR 并保存最好的权重import torch from tqdm import tqdm def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss 0.0 for lr_y, hr_y in tqdm(loader, desctrain, leaveFalse): lr_y lr_y.to(device) hr_y hr_y.to(device) pred model(lr_y) loss criterion(pred, hr_y) optimizer.zero_grad() loss.backward() # 梯度裁剪防偶发爆炸SRCNN 不常炸但小显存大 batch 时有用 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() total_loss loss.item() * lr_y.size(0) return total_loss / len(loader.dataset)criterion用nn.MSELoss()输入和标签都是 0-1 范围别把 0-255 的图直接塞进去。clip_grad_norm_的max_norm1.0是保险丝正常训练很少触发。验证时记得model.eval()和torch.no_grad()否则显存会一路涨。保存权重不要只存state_dict里的网络优化器状态和 epoch 也一起存后面断点续训少受罪torch.save({ epoch: epoch, model: model.state_dict(), optimizer: optimizer.state_dict(), best_psnr: best_psnr }, checkpoints/best.pth)epoch用来续训best_psnr用来判断要不要覆盖文件。只存model.state_dict()也能推理但想接着训练就得重新热身学习率容易把已经收敛的权重打乱。4. 训练与推理全流程从 argparse 到 overlap 切块拼图4.1 训练脚本参数与命令行入口训练脚本用argparse把常用参数暴露出来别把patch_size、scale、lr写死在代码里。下面这段是入口部分的骨架参数名和后面的命令对应import argparse def parse_args(): parser argparse.ArgumentParser() parser.add_argument(--data_dir, typestr, defaultdata/HR) parser.add_argument(--scale, typeint, default3, choices[2, 3, 4]) parser.add_argument(--patch_size, typeint, default96) parser.add_argument(--batch_size, typeint, default16) parser.add_argument(--epochs, typeint, default100) parser.add_argument(--lr, typefloat, default1e-4) parser.add_argument(--seed, typeint, default42) return parser.parse_args()scale控制放大倍数训练和推理必须一致patch_size是 HR 侧裁剪边长batch_size在 8G 显存上先给 16爆了就减半lr从1e-4起步loss 震荡明显就降到5e-5。固定seed方便复现别每次跑完发现结果不一样就怀疑模型。启动命令可以写成python train.py --data_dir data/HR --scale 3 --patch_size 96 --batch_size 16 --epochs 100 --lr 1e-4第一次跑建议把epochs改成 5确认checkpoints/下有文件、验证 PSNR 不是 NaN再开长训练。训练日志里重点看两个数训练 MSE 是否稳定下降验证 PSNR 是否在前几个 epoch 上升后趋于平缓。如果训练 MSE 降、验证 PSNR 不升多半是验证集和训练集分布差太多或者验证时忘了model.eval()。4.2 整图推理的切块与拼接整图直接进网络大图很容易爆显存而且边缘补零会产生接缝。常见做法是切块推理块之间留 overlap拼的时候只取每块中心区域把边界像素丢掉。下面这个函数处理的是已经上采样到目标尺寸的 Y 图切块大小patch256、重叠overlap16可以按显存调import torch from PIL import Image from torchvision.transforms.functional import to_tensor, to_pil_image torch.no_grad() def infer_y(model, y_up, device, patch256, overlap16): model.eval() w, h y_up.size # pad 到 patch 的整数倍避免最后一块越界 pad_w (patch - w % patch) % patch pad_h (patch - h % patch) % patch y_pad Image.new(L, (w pad_w, h pad_h), 0) y_pad.paste(y_up, (0, 0)) out Image.new(L, y_pad.size, 0) step patch - overlap for top in range(0, y_pad.size[1], step): for left in range(0, y_pad.size[0], step): box (left, top, min(left patch, y_pad.size[0]), min(top patch, y_pad.size[1])) block y_pad.crop(box) block_t to_tensor(block).unsqueeze(0).to(device) pred model(block_t).clamp(0, 1).cpu() pred_img to_pil_image(pred.squeeze(0)) # 只取中心区域边缘 overlap 丢掉 inner (overlap // 2, overlap // 2, pred_img.size[0] - overlap // 2, pred_img.size[1] - overlap // 2) out.paste(pred_img.crop(inner), (left inner[0], top inner[1])) return out.crop((0, 0, w, h))patch是单块边长overlap是相邻块重叠像素。step patch - overlap控制滑动步长inner把每块四周各裁掉overlap//2只保留中心预测。拼完后按原图尺寸裁回去。这个写法对 SRCNN 够用缺点是 overlap 区域算了两遍速度略慢如果追求速度可以改成权重图加权平均但代码复杂度会上来。推理时先把 LR 图转 YCbCrY 通道双三次上采样到目标尺寸再走infer_y最后和放大的 CbCr 合并转 RGB 保存。4.3 PSNR 评估与最佳权重保存PSNR 用 MSE 算输入范围必须是 0-1max_val1.0。评估只在 Y 通道做和训练目标一致。下面这个函数放在utils/metrics.py里训练和推理都能复用import torch def psnr(pred, target, max_val1.0): pred pred.clamp(0, max_val) target target.clamp(0, max_val) mse torch.mean((pred - target) ** 2).item() if mse 0: return float(inf) return 10 * torch.log10(torch.tensor(max_val ** 2 / mse)).item()pred和target是[B,1,H,W]或[1,1,H,W]先clamp再算避免模型输出越界把 PSNR 抬得虚高。验证循环里累加每个 batch 的 PSNR取平均后和best_psnr比较大于就保存。别用训练集算 PSNR 来选模型训练集指标只会越来越好选出来的权重多半过拟合。验证集小的话PSNR 波动会比较大可以每 5 个 epoch 评估一次减少偶然性。5. 避坑与常见问题SRCNN 复现时最容易翻车的 5 个点5.1 数据与色彩空间类问题第一个坑RGB 三通道直接训练PSNR 上不去。现象是 loss 能降但放大后的图偏色、边缘发灰和经典 SRCNN 的结果对不上。原因是 SRCNN 的经典设定是在 YCbCr 的 Y 通道上学习CbCr 用双三次插值放大后合并RGB 三通道一起训网络容量被颜色信息分散而且评估口径也变了。解决方式是在 Dataset 里convert(YCbCr)只取 Y 做输入和标签推理时把预测的 Y 和放大的 Cb、Cr 合并再转 RGB。评估 PSNR 也只在 Y 通道算别拿 RGB 三个通道平均。第二个坑下采样和上采样顺序搞反。现象是训练 loss 很快降到很低但推理图几乎没变化像是网络学了个恒等映射。原因是你把 HR 先上采样再下采样或者把 LR 直接当标签网络输入和标签太接近。正确顺序是 HR Y → 双三次下采样到 LR Y → 双三次上采样回 HR 尺寸作为输入HR Y 作为标签。中间那次下采样不能省否则网络没东西可学。检查方法很简单打印一个 batch 的输入和标签看两者是不是明显一个模糊一个清晰。第三个坑卷积 padding 给错输出尺寸对不上。现象是 loss 报 shape mismatch或者输出比标签小一圈。原因是 9×9 卷积不补 padding 会缩 8 像素5×5 会缩 4 像素三层下来尺寸差很多。解决方式是第一层padding4后两层padding2保证输入输出同尺寸。改网络结构时先拿一个torch.randn(1,1,96,96)过一遍打印每层输出形状比肉眼算靠谱。5.2 训练与推理类问题第四个坑整图推理爆显存或者拼出棋盘接缝。现象是验证集小图正常一换大图就 OOM或者拼完的图在块边界有一条亮线。原因是整图卷积的中间特征太大或者切块后直接硬拼边缘补零的预测被保留下来。解决方式是切块推理块之间留 overlap拼的时候只取每块中心区域patch从 256 开始显存不够降到 128overlap给 16 或 32。拼完再检查四个角有没有黑边有黑边说明 pad 或裁剪坐标算错了。第五个坑PSNR 计算范围不对指标虚高或虚低。现象是训练日志里 PSNR 到了 40 多但肉眼看着还是很糊或者 PSNR 只有十几模型明明已经收敛。原因是输入是 0-1PSNR 却按 255 算或者忘了clamp模型输出负值把 MSE 拉大。解决方式是统一用 0-1 范围max_val1.0算之前clamp(0,1)。另外别在 RGB 上算完再和论文的 Y 通道指标比口径不一样比出来没意义。注意这五个点里数据顺序和色彩空间最容易翻车而且症状都像“模型不行”。先打印一个 batch 的输入、标签、输出确认形状和范围再怀疑网络结构和学习率。6. 二次开发把 SRCNN 改成可复用的超分小工具SRCNN 本身容量小直接拿来做产品不够打但拿它当二次开发底座很合适。我通常先加残差连接让网络只学高频输入双三次上采样后的 Y输出x residual训练初期 loss 更稳。改动只在forwarddef forward(self, x): # 残差学习网络学高频最后加回双三次上采样结果 return x self.net(x)self.net最后一层初始化可以改小一点比如乘以 0.1避免一开始残差太大。损失函数可以从 MSE 换 L1视觉上边缘更干净但 PSNR 不一定更高按目标选。批量推理可以在inference.py里加--input_dir和--output_dir复用infer_y别为每张图重启进程。改造点改哪里验证方式残差学习models/srcnn.py的forward前 10 epoch 验证 PSNR 是否更稳L1 损失train.py的criterion对比边缘视觉和 PSNR批量推理inference.py加目录参数跑 20 张图看输出命名和耗时早停验证循环里记best_psnr连续 10 轮不升就停验证改造有没有效我习惯先拿 20 张图做单批次过拟合关掉数据增强固定同一批输入训练 50 个 step看 PSNR 能不能明显上去。如果连 20 张图都过拟合不了说明链路有问题别急着上全量。从那以后我每次复现超分模型都先走一遍“小数据过拟合 单图推理 PSNR 对齐”这三步确认没问题再开长训练。希望帮到你。本文还有配套的精品资源点击获取
