简介:面向深度学习图像超分研究人员与开发者,这份压缩包提供了RDN(残差密集网络)的Pytorch完整复现,代码注释详尽、结构清晰。从数据集制作到模型训练、单图测试、基准验证均有对应脚本,适合正在复现超分算法或需要高倍率重建基准模型的读者。包内共52个文件,以9个Python源码为核心,配合4个模型权重pth文件、png/bmp测试样例、xml工程配置、csv训练日志与README,整体大小327.35MB;目录将data/datasets/epoch/Plt等模块划分明确,便于按需检索,另有PyCharm工程文件可快速打开。目前已有311人学习下载。RDN实现采用模块化设计,颜色空间转换与PSNR/SSIM计算、h5格式训练集制作、多个基准集自动评估等环节均提供独立脚本,可直接复用;x2、x3、x4最优权重按倍率存放,加载即可用于测试或迁移,示例图片的超分结果可直观对比效果,配合配套博文能快速跑通完整流程。
1. 图像超分辨率RDN的PyTorch复现,难的不是网络而是指标与权重对齐
在SwinIR、HAT这类基于Transformer的模型已经把PSNR推到新高度的今天,RDN(Residual Dense Network)仍然是我在重排任务里反复对照的基线:结构对称、训练稳定、上采样端好替换,而且作为CVPR 2018的工作,它对算力的要求远低于后续的注意力模型。这个标题里真正有价值的信息不是“又要复现一个老网络”,而是“注释详细、易读易复用”加上“x2/x3/x4各有一套最优SSIM和PSNR的权重”。换过几个超分项目就会明白:跑通一个RDN训练只需要几十行PyTorch,但想让三个尺度的权重在验证集上都能稳定复现论文指标,功夫全在数据管线、评价口径和上采样端的细节上。这也是我想按工程习惯把RDN重新拆一遍的原因,适合刚入门超分、想要一份能改能跑的参考实现的人,也适合已经在用EDSR或SRResNet、想换个结构做对比实验的人。
2. RDN架构的来路与设计:残差密集块凭什么稳
RDN全称Residual Dense Network,核心动机是同时拿到残差学习的易优化特性和密集连接的复用特性。在它之前,SRResNet和EDSR走的是“宽通道+深残差”路线,特征在经过几十个卷积后,浅层细节大概率衰减;而DenseNet证明密集连接可以缓解梯度消失,但直接把它搬到超分上会导致特征图数量线性膨胀,计算量和显存都扛不住。RDN给出的解是“局部密集、全局稀疏”:在残差块内部做密集连接,块与块之间用普通残差传递,最后在全局做一次特征融合。看上去像两个成功方案的缝合,实际上它解决了超分网络一个容易被忽略的问题——低频信息和高频纹理需要的感受野与通道侧重不同,密集连接让每一层都能看到原始输入,等于在每个块内部做了一次隐式的高频筛选。
2.1 残差密集连接:RDB的密度与局部融合
RDN的基本单元叫RDB(Residual Dense Block),一个RDB内部是若干卷积层,每层输入是“该块初始输入 + 前几层输出”的拼接,层数通常取3或4,卷积输出通道数叫增长率G,论文里常见G=32,主通道数C=64。块的最后,用一层1x1卷积把所有拼接起来的特征图压回C通道,再做一次残差相加。用PyTorch写一个不含激活函数细节的RDB大致如下:
import torch import torch.nn as nn import torch.nn.functional as F class RDB(nn.Module): def __init__(self, nf=64, gc=32, n_convs=4, dilation=1): super().__init__() self.convs = nn.ModuleList() for i in range(n_convs): in_channels = nf + i * gc self.convs.append(nn.Conv2d(in_channels, gc, 3, padding=dilation, dilation=dilation)) self.lff = nn.Conv2d(nf + n_convs * gc, nf, 1) def forward(self, x): features = [x] for conv in self.convs: out = F.relu(conv(torch.cat(features, dim=1))) features.append(out) out = self.lff(torch.cat(features, dim=1)) return out + x代码里需要注意三点:第一,n_convs控制块内卷积层数,增加它不等于加深网络,而是让每一层能看到更早的卷积输出;第二,gc是增长通道数,不是主通道数,改它比改nf更容易在尽量不影响主干的前提下调整参数量;第三,dilation在RDN原论文里没有用到,但这个参数是我做多尺度超分时特意留的,后续如果换成空洞卷积方便直接传值。局部特征融合层的1x1卷积是整个RDB参数量的大头,如果显存吃紧,优先把gc从32改成24,效果比减少n_convs更平滑。
2.2 全局特征融合与全局残差学习:跨块信息是怎样汇总的
单个RDB只能保证块内特征复用,块与块之间如果只是顺序传递,前几个块提取的边缘信息到后端仍可能被稀释。RDN在全部RDB后面加了一个全局特征融合模块,把所有块的输出以及第一个卷积的输出在通道维度上拼接,再用1x1卷积压回C通道。这相当于网络在出口处做了一次“全量特征投票”。全局残差学习则把第一个卷积的输出直接加到融合结果上,让主干通过一条捷径跨越全部RDB,训练初期梯度能顺畅回流。
这里有一个复现时需要点破的细节:很多人会把全局残差连接接到最后一个RDB的输出上,这并不符合原始结构。RDN原文是把F_0(第一个卷积输出)与每个RDB的输出一起送入全局融合,最后再执行F_GF = F_DF + F_0。这样做的意义在于,即使某个RDB内部的密集连接退化成普通卷积,网络依然保留从输入到输出的基础通路。我在第一次复现时曾漏掉F_0参与拼接,训练曲线前50个epoch几乎看不出差别,但最终PSNR低了约0.1dB,这种差距在发表对比实验时是致命的。
| 结构 | 块内连接 | 全局连接 | 训练稳定性 | 适用场景 |
|---|---|---|---|---|
| EDSR | 单流残差 | 仅末端残差 | 依赖大batch与大学习率 | 通道宽、结构简单 |
| RDN | 密集连接+局部融合 | 全局融合+全局残差 | 对学习率更宽容 | 通用超分,中等参数量 |
| RCAN | 通道注意力+残差组 | 组间残差 | 需要较长时间收敛 | 追求极限PSNR且算力充足 |
这个表格是结构层面的横向对比,实际选型还要看退化场景。RDN在压缩伪影去除、真实图像超分这类非理想退化任务中,密集连接带来的鲁棒性往往比EDSR更明显,这是它至今仍被当作强基线的原因。
2.3 上采样末端:x2/x3/x4共享主干,差异只在尾部
RDN对三种尺度的处理方式与EDSR一致:主干特征提取部分完全一样,只在网络末端接一个专门为尺度设计的子网络。常见做法是把最后一个卷积输出的C通道特征送入一层卷积,把通道数变成C * scale^2,再用PixelShuffle重组为高分辨率图像。对x4,可以串联两个x2的像素重组模块;对x3,则需要配置一个通道数为C * 9的卷积。这样设计有一个连带好处:训练好的x2主干可以拿来初始化x4模型,只把上采样端从头训练。实际加速效果取决于数据量和epoch数,小数据集上我一般还是老老实实从零训x4,避免上采样端的随机初始化拖累主干收敛。
3. PyTorch复现RDN:最小可用代码和想清楚再抄的参数
一份“易读易复用”的复现代码,收益最大的部分是网络定义文件。我不会把训练、数据、模型全塞进一个py文件里,而是拆成model.py、dataset.py、train.py、test.py四份,再把所有超参数集中到一个config.py。这样做的原因是,超分辨率实验从x2换到x3或x4,往往只改几个参数,如果参数散落在代码各处,换尺度时会留下隐患。网上很多RDN复现看起来又长又乱,通常不是网络本身复杂,而是作者把每个模块的每个变量都用不同命名方式写了一遍。
3.1 在PyTorch中搭出RDN主干:像素重组与多尺度分支
上采样端的实现决定了这个RDN是否方便切换x2/x3/x4。我习惯把upsampler做成一个独立函数,用scale和nf动态拼装网络结构,而不是为每个尺度单独写一个类。
def make_upsampler(scale, nf): layers = [] if scale in (2, 4): layers += [nn.Conv2d(nf, nf * 4, 3, padding=1), nn.PixelShuffle(2)] if scale == 4: layers += [nn.Conv2d(nf, nf * 4, 3, padding=1), nn.PixelShuffle(2)] elif scale in (3, ): layers += [nn.Conv2d(nf, nf * 9, 3, padding=1), nn.PixelShuffle(3)] else: raise NotImplementedError(f'Scale {scale} is not supported.') return nn.Sequential(*layers)这里有个值得说明的工程选择:x4并没有直接用一个PixelShuffle(4),而是拆成两个x2模块。原因在于,PixelShuffle(4)对应的卷积参数量和计算量是x2模块的4倍,而两个串行的x2模块能增加上采样前的感受野,收敛后的指标通常更好。写完上采样端,整个RDN前向流程就清晰了:
class RDN(nn.Module): def __init__(self, num_blocks=16, nf=64, gc=32, scale=4): super().__init__() self.conv_first = nn.Conv2d(3, nf, 3, padding=1) self.rdbs = nn.ModuleList([RDB(nf, gc) for _ in range(num_blocks)]) self.gff = nn.Conv2d((num_blocks + 1) * nf, nf, 1) self.conv_body = nn.Conv2d(nf, nf, 3, padding=1) self.upsampler = make_upsampler(scale, nf) self.conv_last = nn.Conv2d(nf, 3, 3, padding=1) def forward(self, x): f0 = F.relu(self.conv_first(x)) feats = [f0] f = f0 for rdb in self.rdbs: f = rdb(f) feats.append(f) body = self.conv_body(self.gff(torch.cat(feats, dim=1))) out = f0 + body out = self.upsampler(out) return self.conv_last(out)前向逻辑里最容易被忽略的是feats列表的长度。因为有num_blocks个RDB,加上初始卷积输出f0,拼接后共有num_blocks + 1份特征图,gff层的输入通道必须写对。很多人复现时在这里少算一个f0,训练初期loss能正常下降,但中后期PSNR会比正确实现低0.2dB以上,并且不容易察觉。
3.2 注释详细不等于注释密集:让参数变化成为自解释代码
“注释详细”在我的工程标准里不是每行都写中文注释,而是模块结构清晰、关键设计决策有注释说明。RDB类里的lff、RDN类里的num_blocks + 1,这些地方一旦写错,改起来成本极高。我倾向于在类定义前用三引号写一段“这个类与论文结构的对应关系”,而不是逐行注释卷积的padding是1。对复用而言,更重要的是把超参数的取值范围写在配置里。
class Config: scale = 4 num_blocks = 16 # 论文中的D nf = 64 # 主通道数C gc = 32 # 增长率G patch_size = 64 # 训练时HR裁剪尺寸,x4时LR为16 batch_size = 16 lr = 1e-4 epochs = 600 data_dir = './dataset/DIV2K' log_freq = 500 # 每多少步打印一次训练信息注意patch_size的含义是HR尺寸,LR尺寸由patch_size // scale决定。x4模型用64x64的HR patch,实际进入网络的LR只有16x16,这是一个很多新手会犯迷糊的点。如果显存不够,减小batch_size比减小patch_size的影响更可控,因为patch太小会导致梯度噪声增大。
3.3 环境与依赖版本的隐性约束
PyTorch跑RDN并不需要特殊依赖,但PyTorch版本会影响随机性和卷积行为。以x4权重文件的复现为例,PyTorch从1.x升到2.x之后,cuDNN的卷积算法选择策略变了,同一份代码在1.13和2.4下训练得到的最好权重不会完全一样。如果你要求的是一个能在自己环境里跑出论文水平指标的复现,而不是逐位复现某份权重,那么安装PyTorch时注意CUDA版本与cuDNN的匹配即可,比如Python 3.10配CUDA 12.1的组合包,推理时GPU的NVCC版本和PyTorch的CUDA runtime不必强求一致,但训练时不一致会导致偶发的非法内存访问。最省事的做法是先在CPU上跑一个toy参数训练验证forward/backward完整,再切回GPU训练。
4. 训练到权重产出:数据管线、SSIM/PSNR与x2/x3/x4的选择差异
训练RDN的最终产物是三份权重文件,通常命名为RDN_x2_best.pth、RDN_x3_best.pth、RDN_x4_best.pth,选择依据是验证集上的PSNR或SSIM。这里要先定义清楚“最优”的口径:验证集PSNR最高的权重和验证集SSIM最高的权重往往不是同一份。到底用哪份做发布权重,取决于下游任务更在意像素一致性还是结构相似度。做图像压缩的客户端,我会优先保留SSIM最高;做医学影像重建,PSNR高的权重在评测中更占优势。
4.1 数据管线:HR裁剪、双三次下采样与增强顺序
超分训练集的标准做法是准备高清图像,运行时在内存中裁剪出固定大小的HR patch,再用F.interpolate(..., mode='bicubic', align_corners=False)缩成LR。这个顺序不能反过来,否则LR与HR的对齐方式会直接影响像素损失的计算。常见做法是:
class SRDataset(torch.utils.data.Dataset): def __init__(self, paths, scale, patch_size): self.paths = paths self.scale = scale self.patch_size = patch_size def __getitem__(self, idx): hr = self.load_image(self.paths[idx]) ih, iw = hr.shape[2:] ph, pw = self.patch_size, self.patch_size sh, sw = torch.randint(0, ih - ph, (1,)).item(), torch.randint(0, iw - pw, (1,)).item() hr_patch = hr[:, :, sh:sh + ph, sw:sw + pw] lr_patch = F.interpolate(hr_patch, scale_factor=1/self.scale, mode='bicubic', align_corners=False) return lr_patch, hr_patch代码中F.interpolate缩小时的默认抗锯齿行为与MATLAB的imresize不一致,这是复现论文指标差异最大的来源之一。RDN原实验用的是Matlab的imresize退化核,因此在PyTorch里最接近的替代是torchvision.transforms.functional.resize配antialias=True。数据增强的顺序也有讲究:先做随机旋转和翻转,再做下采样。如果先下采样再翻转,高频信息和低频信息的相对分布会被改变,最后训练出来的模型对翻转后的LR输入不友好。
4.2 训练参数:L1还是Charbonnier,学习率要不要warmup
RDN原文用L1损失,L1在超分任务中比L2能保留更多高频细节,收敛也更稳。训练参数表格如下:
| 参数 | 建议值 | 理由 |
|---|---|---|
| 损失函数 | L1 (MAE) | 对离群像素不敏感,纹理更锐利 |
| 优化器 | Adam,betas=(0.9, 0.999), eps=1e-8 | 超分任务不用SGD,收敛太慢 |
| 初始学习率 | 1e-4 | 高于这个值容易在RDB密集连接处震荡 |
| 学习率调整 | 每200个epoch乘0.5 | 与step-based调参配合,稳定性好 |
| 梯度裁剪 | clip_grad_norm_(max_norm=1.0) | 防止密集连接带来的梯度爆炸 |
| batch_size | 16(x2)/ 12(x4) | x4的LR更小,batch可以略大但不建议过大 |
Charbonnier损失是L1的可导近似,即sqrt(x^2 + eps^2),通常eps取1e-3到1e-6。实际训练中L1已经够用,Charbonnier的优势主要体现在训练后期收敛更平稳,但会牺牲一点速度。如果只做复现,直接用L1并把学习率warmup到1e-4再衰减,能在200个epoch内达到论文指标的98%以上。
权重文件的保存逻辑值得留意,不要简单地保存最近一个epoch,而应该在验证集上逐epoch计算PSNR,并只在超过历史最优时覆盖权重。训练一个x2模型在中等显存卡上通常要一整天,如果只保存最后一份权重,可能错过中后期指标最高的那个时间点,这也是标题里“最优SSIM和PSNR的模型权重文件”真正的含义。
4.3 为什么x2、x3、x4要分别训练,而不是共用一份x4权重插值缩放
超分模型不像普通图像缩放,x2模型的输出直接resize到x4,效果远差于原生x4模型。原因有二:其一,x2模型只学过2倍下采样的退化分布,对4倍退化引入的高频信息丢失无能为力;其二,上采样端的结构完全不同,x2的PixelShuffle输出通道是nf*4,x3是nf*9,x4是两次nf*4,权重文件本身的结构就不兼容。所以三份权重必须独立训练、独立选择最优指标。如果算力有限,可以在x2模型训练完成后,复制其主干权重作为x4的初始化,只重新训练上采样端,这比从零训练快大约15%的epoch。
4.4 SSIM与PSNR计算口径:边界裁剪和Y通道是硬性规定
评价超分结果时,PSNR和SSIM的计算不是把两张RGB图直接扔进torchvision.utils那么简单。公认的计算口径是先转YCbCr,只取Y通道,并把图像四周按尺度大小裁剪掉对应像素,再逐像素计算。裁剪的原因是下采样和卷积填充会在边缘产生不可靠的像素,这部分会拉低PSNR。
import torch import numpy as np def convert_y(img): # img: RGB图像,float范围0-255 img = img.permute(1, 2, 0).numpy() img = img.astype(np.float64) y = 16.0 + (64.738 * img[..., 0] + 129.057 * img[..., 1] + 25.064 * img[..., 2]) / 255.0 return y def calc_psnr(sr, hr, scale, border=0): sr, hr = sr[border:-border, border:-border], hr[border:-border, border:-border] mse = np.mean((sr - hr) ** 2) if mse == 0: return float('inf') return 10 * np.log10(255.0 * 255.0 / mse)SSIM建议直接使用skimage.metrics.structural_similarity,但要把data_range设为255,win_size根据图像尺寸调整,并且同样只算Y通道。很多公开权重文件的PSNR数值与论文不一致,基本都是边界裁切或者通道选择问题,而不是模型训练问题。这个口径从复现阶段就固化下来,才能保证“最优权重”的评价标准是可复现的。
5. 用权重文件做推理与验证:checkpoint加载、评价脚本和最容易踩的坑
拿到x2/x3/x4权重文件后,第一步不是跑整张图看效果,而是核对权重结构。RDN的权重文件里通常是model_state_dict、optimizer_state_dict和一个自定义的metrics字段,存放训练时的验证PSNR。加载时注意map_location指定设备,以及在PyTorch 2.0之后load_state_dict对严格模式默认开启,键名不匹配会直接报错。
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') state = torch.load('RDN_x4_best.pth', map_location='cpu') print(state.keys()) # 先看里层结构 print(state['metrics']) # 打印训练时记录的验证PSNR/SSIM model = RDN(scale=4).to(device) model.load_state_dict(state['model']) model.eval()一个常见的坑是训练时用了nn.DataParallel,权重文件的模型键名全部带上了module.前缀。解决方法是加载后做一次strip:state['model'] = {k.replace('module.', ''): v for k, v in state['model'].items()},否则就会报大小不匹配或缺少键。另一个坑是权重文件由PyTorch旧版本保存,新版本加载时如果遇到weights_only相关提示,在确认来源可信的前提下可把torch.load的weights_only设为False。
推理时尽量用torch.inference_mode()而不是torch.no_grad(),前者会关闭一部分与自动微分无关的追踪逻辑,在V100及以上的GPU上能快10%左右。输出图像前先把模型输出限制到[0, 1]或[0, 255]再转uint8,常见做法是torch.clamp_(pred, 0, 1)后乘255,不要依赖卷积层末尾的clip,卷积输出越界是常态。
验证一份权重好坏的最快方式,是用Set5或Set14的标准图像,把HR图和一张双三次放大后的LR图同时送入评价脚本,对比calc_psnr的结果是否落在该模型预期范围。如果结果与记录值偏差超过0.5dB,优先检查是否存在翻转或通道顺序问题。实际推理工程里,我还建议把三份权重放在同一目录,通过命令行参数传scale自动选择对应文件,避免混淆x2和x4的模型,这个看似简单的路径错误在自动化流程里会造成难以察觉的指标下滑。
本文还有配套的精品资源,点击获取