
简介图像超分辨率WDSR的PyTorch复现项目面向深度学习与超分重建方向的开发者、研究人员适用于算法复现、模型对比及论文实验场景。项目对WDSR-A和WDSR-B两种结构均做了完整实现训练、验证、测试代码齐备注释详细并配有科研常用的Loss与PSNR曲线绘制脚本方便直接嵌入自己的实验流程。资源共47个文件以16个Python脚本为核心涵盖数据预处理DIV2K转h5、模型定义、参数配置、训练评估等环节另有12个tar格式模型权重文件覆盖x2、x3、x4倍率下最优与最新检查点其余为配置文件与说明文档压缩包整体约95.87MB。数据划分、数据增强、SSIM计算等模块单独封装结构清晰便于按需修改和二次开发。已有24770人学习下载。对于想快速上手超分辨率任务、复现WDSR或基于预训练权重开展对比实验的读者这份资源提供了从数据处理到结果可视化的完整闭环值得收藏。1. 为什么是WDSR一个在PSNR之外把“宽度”玩明白的结构如果你在超分方向调过EDSR或SRResNet应该会注意到一个老问题把网络做深做宽PSNR是涨上去了但参数量和显存也跟着失控。WDSRWide Activation for Super Resolution论文里给了另一个思路——不需要一味加深度而是在ReLU激活前用线性低秩卷积把特征通道撑宽激活后再压缩回低维从而在相同参数量下获得更高的PSNR和更清晰的纹理。这个仓库把WDSR-A和WDSR-B两种变体用PyTorch完整复现带x2、x3、x4的best/latest权重文件还包含训练、验证、benchmark测试、SSIM计算和科研绘图脚本。适合刚接触图像超分辨率重建的入门者也适合需要快速拿现成权重跑对比实验的老手——你不需要从零训练load权重就能在自己的图上复现超分效果。2. 数据管线从DIV2K到h5及数据增强的细节超分实验的复现性很大程度取决于数据预处理。DIV2K是超分benchmark的默认训练集但这个仓库没有直接用原始PNG流式读入而是先用core/data/div2k.py把训练集和测试集打包成h5格式再转成Tensor。这样做的直接好处是训练时随机裁剪、翻转、旋转都在内存中的小patch上进行不必反复IO大图训练速度快很多而且h5文件方便跨机器拷贝。2.1 图像读取与Tensor转换core/data/utils.py里的基础操作utils.py里集中了读图、PIL转Numpy、数据增强等操作。下面这段代码是基于该项目常见写法整理出的核心逻辑和仓库里的utils.py行为一致import numpy as np from PIL import Image def load_img(filepath): 读取图像并转为RGB float数组值域[0,1] img Image.open(filepath).convert(RGB) return np.array(img).astype(np.float32) / 255.0 def rgb2ycbcr(img): RGB转YCbCr超分常用只训练/评估Y通道 if img.ndim 3: img img.astype(np.float64) y 16.0 (64.738 * img[..., 0] 129.057 * img[..., 1] 25.064 * img[..., 2]) / 256.0 return y def augment(lr, hr, hflipTrue, rotTrue): 数据增强随机水平翻转和90度旋转 if hflip and np.random.rand() 0.5: lr lr[:, ::-1, :] hr hr[:, ::-1, :] if rot: k np.random.randint(0, 4) lr np.rot90(lr, k) hr np.rot90(hr, k) return lr.copy(), hr.copy()这里load_img先做归一化到[0,1]避免了后续模型输出层再乘255带来的数值不稳定。rgb2ycbcr返回的Y通道是uint8转换后的float范围在[16,235]之间计算PSNR时通常只取Y通道且裁掉边界。augment里的np.rot90传入k0时不会复制所以要.copy()防止后面h5写入时共享内存出问题。值得注意的是这个仓库的DIV2K专属均值偏移放在core/model/common.py里训练时先对输入减均值输出再加回这一步在测试图像时别忘了。2.2 制作h5数据集div2k.py为什么要把图片patch化div2k.py负责把DIV2K的HR原图切patch同时生成对应的LR。常见做法是先用bicubic下采样得到LR然后随机或滑动取patch对。整个流程类似import h5py import numpy as np from PIL import Image def make_h5(hr_dir, lr_dir, h5_path, patch_size192, stride96, scale4): hr_list sorted(hr_dir.glob(*.png)) with h5py.File(h5_path, w) as f: for i, hr_path in enumerate(hr_list): hr load_img(hr_path) lr_path lr_dir / hr_path.name.replace(x4, fx{scale}) # 假设LR已存在 lr load_img(lr_path) # 滑动裁剪HR和对应LR for top in range(0, hr.shape[0] - patch_size 1, stride): for left in range(0, hr.shape[1] - patch_size 1, stride): hr_patch hr[top:top patch_size, left:left patch_size] lr_top top // scale lr_left left // scale lr_patch lr[lr_top:lr_top patch_size // scale, lr_left:lr_left patch_size // scale] # 增强后写入 lr_patch, hr_patch augment(lr_patch, hr_patch) f.create_dataset(fhr_{i}_{top}_{left}, datahr_patch) f.create_dataset(flr_{i}_{top}_{left}, datalr_patch)这个脚本的patch_size和stride直接决定训练样本数量。patch_size192在DIV2K上能得到较大的感受野但显存有限时建议降到128stride可以小于patch_size让patch之间重叠但对训练效果提升有限反而增加训练集冗余。实际训练时如果能在线随机裁剪就没有必要离线存h5。这个仓库选择h5是因为训练前已经固定好patch复现论文指标时不需要再抖动随机性。3. WDSR模型架构与实验配置A/B变体如何选WDSR的核心贡献是提出了一种“宽激活”残差块。传统的SRResNet先降维再卷积再升维ReLU放在低维空间上而WDSR反过来——先用1x1卷积把通道数线性放大比如从64放大到256再做ReLU然后通过低秩卷积或1x1卷积把通道压缩回64。这个看似简单的顺序调整让信息在激活前保持高维非线性变换后不丢失边缘细节理论上能加速训练收敛也让相同参数量下的表示能力更强。3.1 WDSR-A与WDSR-B的差异WDSR-A采用最直接的宽激活残差单元输入先经过一个1x1卷积把通道数从f扩展到f*r接着ReLU再用一个1x1卷积把通道压回f最后加一个可选的3x3卷积。WDSR-B则在这个基础上拆成两个并行分支其中一个分支只做恒等映射另一个分支做低秩卷积最后相加。用PyTorch描述的话WDSR-B的残差块大概长这样import torch.nn as nn class WDSRBlockB(nn.Module): def __init__(self, n_feats64, expansion4, wnFalse): super().__init__() self.expand nn.Conv2d(n_feats, n_feats * expansion, 3, padding1) self.contract nn.Conv2d(n_feats * expansion, n_feats, 1, padding0) self.relu nn.ReLU(inplaceTrue) if wn: self.expand nn.utils.weight_norm(self.expand) self.contract nn.utils.weight_norm(self.contract) def forward(self, x): y self.relu(self.expand(x)) y self.contract(y) return x yn_feats对应权重文件名里的f32expansion对应文件名里的r4也就是宽度倍率。WDSR-B用3x3卷积做通道扩展感受野更大但参数量略高于A。实际训练时如果感觉A收敛慢可以换成B试如果显存不大A的f32r4默认配置是性价比最高的。3.2 权重文件名的含义与option.py参数对应仓库里的权重文件名如WDSR-A-f32-b16-r4-x3-best.pth.tar拆解后是模型类型A基本特征数32残差块数16宽度倍率4超分倍率3后缀best表示验证集上PSNR最高的权重latest表示最近保存的权重。对应的option.py里核心参数如下参数含义默认值建议modelWDSR-A或WDSR-BWDSR-B小数据集用A大数据集用Bn_resblocks(b16)残差块数量168可提速32可提升但显存翻倍n_feats(f32)基础通道数3264提升明显但显存暴涨n_expand(r4)宽度倍率44是论文最佳6偏重纹理scale超分倍率2/3/4按目标设这些参数在weight文件和option.py里一致性校验如果你自己训练时改了n_feats不要直接加载仓库给的权重否则Pytorch会报state_dict不匹配。调试阶段可以在option.py里把--seed固定保证每次裁图和权重初始化一致。4. 训练闭环train.py、eval.py与benchmark验证拿到这个仓库第一步不是跑测试而是把训练、验证、benchmark三条路径走通。这样你才知道权重文件是怎么来的也方便后续在自己数据集上微调。训练入口是train.py它负责数据加载、模型构建、损失计算、梯度反向和checkpoint保存。4.1 训练脚本结构与执行流程先看命令行怎么启动训练以x3倍率、WDSR-B为例python train.py --model WDSR-B --scale 3 --n_feats 32 --n_resblocks 16 --n_expand 4 --batch_size 16 --lr 1e-4 --epochs 300 --save_dir epoch/训练循环里的关键部分简化自train.pyfrom core.model import WDSR_B from core.data import get_training_set from option import args from torch.utils.data import DataLoader from torch.optim import Adam import torch # 模型初始化加载数据 model WDSR_B(n_resblocksargs.n_resblocks, n_featsargs.n_feats, n_expandargs.n_expand, scaleargs.scale) train_set get_training_set(args.scale) train_loader DataLoader(train_set, batch_sizeargs.batch_size, shuffleTrue, num_workers4, pin_memoryTrue) criterion torch.nn.L1Loss() # 超分常用L1 optimizer Adam(model.parameters(), lrargs.lr) # 训练一个epoch for lr, hr in train_loader: lr, hr lr.cuda(), hr.cuda() sr model(lr) loss criterion(sr, hr) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) # 防止梯度爆炸 optimizer.step()这里用L1Loss而不是L2(MSE)因为L1在超分任务上通常能产出更锐利的边缘PSNR略低但主观更好。clip_grad_norm_是训练早期稳定性的保险丝当lr较大时尤其需要。每训练几个epoch脚本会调用eval.py在DIV2K验证集上算一次PSNR如果比历史最好就保存为*best.pth.tar否则只保存*latest.pth.tar。4.2 eval.py验证单个模型验证脚本的用法非常直接python eval.py --model WDSR-B --scale 4 --n_resblocks 16 --n_feats 32 --n_expand 4 --resume epoch/WDSR-B-f32-b16-r6-x4-best.pth.tar验证时有一个容易忽略的细节必须在模型输出后立即用clamp把像素值裁剪到[0,1]否则某些新结构会产生略微超出范围的输出导致最终PSNR虚高。仓库里常见的验证逻辑是def eval_psnr(sr, hr): sr sr.clamp(0, 1) # 先裁剪 mse ((sr - hr) ** 2).mean() psnr 10 * torch.log10(1.0 / mse) return psnr.item()注意这里的1.0是像素范围。如果归一化后hr在[0,1]就用1如果你用了图像均值偏移记得先把均值加回来再计算否则PSNR会偏低。4.3 test_benchmark.py批量测试5个benchmarktest_benchmark.py是给论文对比用的它一次性跑Set5、Set14、BSD100、Urban100、Manga109五个数据集输出每个数据集的平均PSNR/SSIM。运行方式python test_benchmark.py --model WDSR-B --scale 2 --resume epoch/WDSR-B-f32-b16-r6-x2-best.pth.tar --data benchmark_root/脚本内部会对每张图做多尺度自集成self-ensemble来提高分数把输入翻转四个方向再每个方向测一次最后取平均。这会把推理时间放大8倍但PSNR通常能提升0.1~0.2dB。如果你的项目是需要跑实时推理就不要开self-ensemble直接测原始方向即可。4.4 draw_evaluation.py科研绘图训练完成后draw_evaluation.py负责把日志里的loss和PSNR曲线画出来。仓库在epoch/目录下会保留每个epoch的记录绘图脚本读取这些json或txt后生成两张图。以下是我整理出的绘图核心import matplotlib.pyplot as plt def plot_metrics(history, save_path): plt.figure(figsize(8, 5)) plt.plot(history[epoch], history[loss], labelTraining Loss) plt.xlabel(Epoch) plt.ylabel(Loss) ax2 plt.twinx() ax2.plot(history[epoch], history[psnr], colororange, labelPSNR(dB)) ax2.set_ylabel(PSNR (dB)) plt.legend() plt.tight_layout() plt.savefig(save_path, dpi300)绘图的亮点是双y轴左边loss右边PSNR在论文里能直观看到loss下降的同时PSNR上升。如果要用于正式科研绘图建议把线条加粗、去掉顶部右边框字号改成小五号这样和期刊字体统一。5. 用官方权重跑自己的图像test.py与SSIM计算这个仓库最省事的地方在于你已经有了不同倍率的best权重不需要真正训练就能复现论文级结果。把权重放在epoch/下随便找一张自然图片执行python test.py --input data/mypic.png --output data/mypic_x4.png --model WDSR-B --scale 4 --resume epoch/WDSR-B-f32-b16-r6-x4-best.pth.tartest.py会读取图像转换到YCbCr只对Y通道做超分CbCr直接用bicubic上采样最后合并保存。这样做的原因是WDSR训练时只对Y通道监督CbCr用插值不容易出现明显色偏能省一半计算量。如果你需要RGB端到端超分需要自己写一个把RGB三通道都过网络的脚本。5.1 在自己数据上验证SSIM仓库自带的pytorch_ssim库提供了SSIM计算接口使用前先确认参考图像和超分图像尺寸一致import torch from pytorch_ssim import SSIM from PIL import Image import numpy as np def compute_ssim(hr_path, sr_path): hr load_img(hr_path) # [0,1] sr load_img(sr_path) hr_t torch.from_numpy(hr).permute(2, 0, 1).unsqueeze(0) sr_t torch.from_numpy(sr).permute(2, 0, 1).unsqueeze(0) ssim_module SSIM(window_size11, size_averageTrue) return ssim_module(hr_t, sr_t).item()window_size11是DIV2K上的通用配置如果图像尺寸很小比如小于64像素建议把window_size设成奇数且小于图像尺寸否则会报padding错误。注意这个SSIM是基于高斯权重窗口的和skimage的compare_ssim(win_size11)不一定完全对齐对比指标时最好统一用同一个库。5.2 容易踩的坑边界效应和Y通道指标计算PSNR/SSIM时超分图像边缘往往有meshgrid导致的伪影因此benchmark评测时通常会先去掉scale像素的边界再计算。做法很简单crop scale hr hr[crop:-crop, crop:-crop] sr sr[crop:-crop, crop:-crop]另外如果发现加载仓库权重后PSNR和readme对不上先检查输入图像是否对齐到/scale的整数倍。WDSR的反卷积层会要求输入尺寸对倍率取整否则模型内部上采样后尺寸差1个像素指标掉0.3dB以上。遇到这种情况最常见的做法是用torch.nn.functional.interpolate将输入先放大到合适的尺寸再做超分最后再裁剪回目标尺寸。如果你要在论文里用这些权重出图建议额外保存超分前的低清图像和超分后的结果并标注好哪个权重文件、哪个模型版本方便审稿时复现。这个仓库的latest权重是训练末尾状态best是验证集最优论文里一般用best表格数据latest只有在你需要继续预训练时才用。本文还有配套的精品资源点击获取