简介:在数字图像处理中,将低分辨率图像重建为高分辨率图像一直是计算机视觉的基础技术需求。深度学习尤其是卷积神经网络的兴起,为这项经典任务提供了全新的解决路径。超分辨率模型的核心挑战在于如何让深层网络稳定收敛并有效恢复高频纹理细节。VDSR模型通过残差学习让网络仅预测高低频差异,配合梯度裁剪防止梯度爆炸,再利用多尺度训练统一处理多档放大倍数,仅用20层卷积便在超分任务中取得了显著效果。掌握VDSR的PyTorch实现,不仅有助于理解深度学习的核心训练技巧,也能为图像增强、老照片修复、视频超分等工程实践提供可靠的落地方案。围绕模型结构、训练管线、测试评估与避坑经验展开,帮助从业者快速复现并部署高精度超分工具。
1. 论文复现VDSR:一个2016年的超分模型,为什么现在仍值得用Pytorch跑一遍
把一张老照片放大两倍,边缘全是锯齿、文字糊成一团,这是每个做过图像处理的从业者都撞过墙的日常。论文复现VDSR(Very Deep Super Resolution)解决的就是把低分辨率图像重建成高分辨率图像这个单图超分任务,它的Pytorch实现只有二十层普通卷积和一条残差捷径,结构简单到可以在一张消费级显卡上完成训练。它难在训练技巧而不是模型骨架,也正因为如此,它是复现深度学习论文的绝佳练手对象。适合人群是那些已经会用Pytorch写分类网络、但没跑通过超分这类低层视觉任务的人,还有想把图像增强写进自己采集系统里的工程师。
复现VDSR能从一个具体网络里看到三件有价值的东西:残差学习让深层网络收敛、梯度裁剪让高学习率训练不炸、多尺度训练让一个模型处理多档放大倍数。这三件事在今天很多SR模型里仍在反复出现,看懂了它们,再去看EDSR、SRGAN、ESPCN这些后续模型,会发现很多设计其实是顺着同一个思路长出来的。
2. VDSR原理拆解:20层卷积为什么能赢,残差和梯度裁剪为什么缺一不可
2.1 从SRCNN的三层到VDSR的二十层:深度对单图超分意味着什么
SRCNN是2014年的开山之作,结构只有三层卷积:先把插值放大后的LR图做特征提取,中间一层做非线性映射,最后一层重建出HR图。三层结构带来的直接限制是感受野太小。
感受野的计算公式很简单:连续L层3×3卷积(padding保持尺寸,stride=1),感受野 = 1 + 2×(L-1)。SRCNN三层时约为13×13像素,VDSR二十层是39×39像素。同样一个输出像素要依赖输入图像上多大一块区域来重建,SRCNN只能看到身边13个像素以内的局部纹理,而VDSR能看到接近40像素范围内的结构信息。超分重建最怕的就是局部纹理断裂——一个像素周围看到了更完整的边缘走向,重建出来的线条自然就连续。这也是VDSR敢把网络从三层直接推到二十层的核心动机。
但网络加深不等于白拿收益。二十层的ReLU卷积网络用大的学习率一开训,梯度经过多层反向传播会剧烈放大或消失。VDSR没有用BatchNorm,也没有特殊激活函数,它靠的是残差学习和梯度裁剪这两把钥匙把深度这条路走通了。对比之下,SRCNN之所以停留在三层,正是因为当年没有这两样东西兜底,继续加深只会让训练loss无法下降。
| 维度 | SRCNN | VDSR |
|---|---|---|
| 层数 | 3层 | 20层 |
| 感受野 | 约13×13 | 约39×39 |
| 参数量量级 | 约5.7万 | 约66.6万 |
| 残差学习 | 无 | 有 |
| 梯度裁剪 | 无 | 有,阈值0.4 |
2.2 残差学习才是模型的核心设计:网络只学"补全量"而不是图片本身
先明确一个前提:VDSR的输入不是原始低分辨率图,而是已经用双三次插值放大回目标尺寸的LR上采样图,记作LR_up。这颗图和目标HR图之间的差异,主要集中在边缘、纹理这些高频细节上,大片的平坦区域两者几乎一样。
如果让网络直接从LR_up去学习一整张HR图,它需要再造出全部像素的数值,包括那些本来就已经很接近的部分。残差学习的做法是让网络只预测“补全量”:r(x) = HR - LR_up。对一张自然图像而言,这个量是近似零均值的稀疏信号,90%以上的像素数值落在[-0.05, 0.05]之间,只有边缘处有较大响应。预测一个稀疏信号比预测一个完整图像要简单得多,这也是VDSR在二十层深度下只用MSE损失就能稳定训练的底层原因。
代码里的实现是在forward的返回处把输入加回来:
def forward(self, x): # 网络只预测残差,最后把输入连接加回去 residual = self.body(x) return x + residual这里的x是LR_up,加回来的操作就是所谓的跳过连接。复现时最容易犯的错误是只输出residual而不加x,训练时loss也降得下去,但推理时PSNR会卡在一个低水平上,因为模型给出的只是残差,不是最终图像。这个坑在第4章还会单独展开。
2.3 多尺度训练与梯度裁剪:论文中三个训练设置背后的动机
论文实验部分给出了几个看似简单但极关键的训练设置:学习率初始1e-3,每20个epoch衰减到原来的1/10;梯度裁剪阈值0.4;训练数据同时包含2、3、4倍三种下采样。
学习率1e-3对当时的卷积网络来说并不算小,尤其对二十层结构而言。VDSR敢于用大学习率,依赖的正是残差学习带来的较平滑的损失曲面,再加上梯度裁剪兜底。梯度裁剪相当于在每个批次里检查所有参数的梯度范数,一旦超过阈值就给梯度按比例缩回去,防止某一步的梯度异常大导致参数跳动过头。论文给出的0.4是一个经验值,我之后尝试过0.1和1.0,0.1让收敛变慢,1.0在数据方差大的时候偶尔会飘,0.4确实是稳妥的选择。
多尺度训练的实现思路并不复杂:每个训练批次随机抽一个scale(2、3、4),对这个scale做对应的下采样与上采样造数据。这样同一个模型在训练过程中就见了三种退化程度的样本,推理时无需知道测试图用了哪个倍数,直接用模型预测即可。这套做法后来被很多SR模型沿用,成为节省训练成本的通用技巧。
2.4 这张网络到底有多少参数、多大的感受野
| 层范围 | 输入通道 | 输出通道 | 卷积核 | 该层参数量 |
|---|---|---|---|---|
| 第1层 | 1 | 64 | 3×3 | 576 + 64 = 640 |
| 第2~19层(18层) | 64 | 64 | 3×3 | 36864 + 64 = 36928 / 层 |
| 第20层 | 64 | 1 | 3×3 | 576 + 1 = 577 |
总计约 640 + 18×36928 + 577 ≈ 66.6 万参数。放在今天看这是一个很轻量的网络,但它的有效感受野是39×39,刚好覆盖了训练patch 41×41的实际计算区域。为什么patch要取41?因为每个输出像素需要看到完整的39×39输入视野,patch再小一点,边缘位置的像素就看不全,训练时学到的映射关系与推理时不一致。
参数少、结构规整,意味着在一张普通显卡上训练速度很快、内存占用小,这也是复现它作为入门练手项目的优势。当然它也有时代局限:没有BatchNorm导致网络对初始化比较敏感,固定20层结构无法按图像复杂度动态调整,这些正是后来EDSR等模型着力解决的问题。对复现而言,这些局限反而提供了一个很好的对照系。
3. 用Pytorch实现VDSR训练管线:环境、数据、模型与训练代码
3.1 用Anaconda搭建Pytorch训练环境(含WSL/AMD GPU情况)
先解决环境问题。VDSR模型只有66万参数,对设备要求极低,环境搭建的重点是让PyTorch正确识别GPU,而不是追求高端显卡。最常见的做法是用Anaconda创建一个独立环境:
conda create -n vdsr python=3.8 -y conda activate vdsr pip install torch torchvision装完先验证GPU能否被PyTorch识别:
import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果torch.cuda.is_available()输出True,说明CUDA已经正常接入。如果你的机器是NVIDIA GPU,注意Windows和WSL2是两套驱动体系:WSL2里不要重复安装Linux驱动,宿主机Windows装好对应CUDA版本驱动,WSL里的PyTorch就能直接看到CUDA。如果你用的是AMD 7900XTX这类A卡,可以安装ROCm版本的PyTorch,安装命令同样是pip,只是wheel来源不同;没有N卡也没有A卡时直接用CPU跑小patch训练也不是不行,VDSR单batch的前向计算量很小,只是会慢一些。
环境这一步最容易翻车的是版本错配,比如驱动只支持CUDA 11.3却装了要求CUDA 11.8的PyTorch wheel,表现就是运行时cuda初始化报错,但系统里nvidia-smi查起来又正常。排查时先看PyTorch官方对CUDA的最低版本要求,再对照驱动版本。跑超分任务不建议一上来就追求最新CUDA,稳定是第一位。
3.2 准备HR/LR训练数据:切patch、双三次上采样、数据增强
VDSR训练时不需要专门的低分辨率数据集,而是从HR图动态生成LR图。常规做法是准备一批高分辨率图像,T91、General-100、DIV2K都是常见来源。训练时随机切patch,对该patch做双三次下采样再上采样,得到LR_up作为网络输入,原始patch作为标签。
推荐目录结构:
data/ ├── train/ │ ├── T91/ │ ├── General100/ │ └── DIV2K_train_HR/ └── test/ ├── Set5/ ├── Set14/ └── BSD100/数据加载器代码可以这样写:
import cv2 import numpy as np import torch from torch.utils.data import Dataset import glob def load_image_paths(folder): # 读入指定目录下所有png/jpg/bmp图片 exts = ['*.png', '*.jpg', '*.bmp'] paths = [] for ext in exts: paths.extend(glob.glob(folder + '/' + ext)) return paths class VDSRDataset(Dataset): """从HR图实时切patch生成LR/HR训练对,支持多尺度采样与数据增强""" def __init__(self, hr_image_paths, patch_size=64, scales=(2, 3, 4)): self.images = [] for path in hr_image_paths: img = cv2.imread(path, cv2.IMREAD_GRAYSCALE) if img is None: continue # 统一转成float32,像素值归一化到[0,1] img = img.astype(np.float32) / 255.0 self.images.append(img) self.patch_size = patch_size self.scales = scales def __len__(self): # 每个epoch自定义样本数,避免epoch过长 return 2048 def __getitem__(self, idx): img = self.images[np.random.randint(len(self.images))] scale = np.random.choice(self.scales) ps = self.patch_size h, w = img.shape # 图像太小时先做padding,防止越界 if h < ps or w < ps: pad_h, pad_w = ps - h + 8, ps - w + 8 img = np.pad(img, ((0, pad_h), (0, pad_w)), mode='edge') h, w = img.shape # 随机裁切HR patch i = np.random.randint(0, h - ps) j = np.random.randint(0, w - ps) hr_patch = img[i:i+ps, j:j+ps] # 先下采样再上采样,得到LR_up lr_size = max(1, ps // scale) lr = cv2.resize(hr_patch, (lr_size, lr_size), interpolation=cv2.INTER_CUBIC) lr_up = cv2.resize(lr, (ps, ps), interpolation=cv2.INTER_CUBIC) # 随机翻转与旋转,增加样本多样性 if np.random.rand() < 0.5: hr_patch, lr_up = np.flip(hr_patch, 0), np.flip(lr_up, 0) if np.random.rand() < 0.5: hr_patch, lr_up = np.flip(hr_patch, 1), np.flip(lr_up, 1) k = np.random.randint(0, 4) if k > 0: hr_patch, lr_up = np.rot90(hr_patch, k), np.rot90(lr_up, k) # 增加channel维度,转成tensor返回 lr_up = torch.from_numpy(lr_up.copy()).unsqueeze(0) hr_patch = torch.from_numpy(hr_patch.copy()).unsqueeze(0) return lr_up, hr_patch这段代码里有三个参数值得留意。第一个是patch_size,64是常用折中值,显存紧张时可以降到48,但不建议低于41,否则感受野打不满。第二个是scales=(2,3,4),对应论文的多尺度训练,如果只想要单一倍数,改成(2,)即可。第三个是len方法返回的2048,这个数字不是必须的,它决定一个epoch处理多少个样本,小数据集下建议保留,避免一个epoch遍历完所有图像后过拟合严重。
cv2.resize这里全部使用INTER_CUBIC双三次插值,是为了与论文的退化假设保持一致。如果测试时用了别的插值方式,训练和测试分布不一致,PSNR会掉得厉害。
3.3 VDSR模型定义:结构、初始化与参数量核对
模型结构严格按照论文描述实现。核心是20层3×3卷积,前19层每个都跟一个ReLU,最后一层直接输出,然后用残差连接把输入加回来。
import torch import torch.nn as nn class VDSR(nn.Module): """VDSR网络结构:20层卷积,64通道,输出为残差""" def __init__(self, num_channels=1, num_features=64, num_layers=20): super(VDSR, self).__init__() layers = [] # 第一层:输入通道扩展到64 layers.append(nn.Conv2d(num_channels, num_features, kernel_size=3, padding=1)) layers.append(nn.ReLU(inplace=True)) # 中间18层:保持64通道不变 for _ in range(num_layers - 2): layers.append(nn.Conv2d(num_features, num_features, kernel_size=3, padding=1)) layers.append(nn.ReLU(inplace=True)) # 最后一层:输出通道还原为1,这里不再接ReLU layers.append(nn.Conv2d(num_features, num_channels, kernel_size=3, padding=1)) self.body = nn.Sequential(*layers) self._initialize_weights() def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): # 使用Kaiming初始化,适配ReLU特征分布 nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu') nn.init.constant_(m.bias, 0) def forward(self, x): residual = self.body(x) return x + residual这里有一个容易漏掉的细节:最后一层卷积后面不接ReLU。原因很简单,残差信号的数值允许为负,边缘两侧一正一负才符合真实残差的分布,ReLU会把负值截断,破坏学习目标。
初始化用了Kaiming初始化,这与论文里的“adjusted initialization”不完全一致,但实测效果差异很小,也是复现中最常见的初始化选择。想核对模型结构是否正确,可以打印参数量:
def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) model = VDSR() print(count_parameters(model)) # 预期约665921输出结果在66万附近基本就说明结构对了。如果差很多,检查是不是把中间层数写成了21或19。
3.4 训练循环:MSE损失、SGD动量、每20个epoch衰减学习率
训练主循环相对固定,直接用MSE作为损失函数,优化器用SGD而不是Adam,这一点和很多图像分类任务不同。VDSR论文使用SGD、momentum=0.9、无weight_decay,配合梯度裁剪,在超分任务上表现稳定。
import torch.optim as optim from torch.utils.data import DataLoader def train_vdsr(model, train_dataset, epochs=50, batch_size=16): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = model.to(device) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True, drop_last=True) # 论文配置:SGD、momentum=0.9 optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9) criterion = nn.MSELoss() start_epoch = 0 for epoch in range(start_epoch, epochs): model.train() epoch_loss = 0.0 # 每20个epoch学习率乘以0.1 if epoch in (20, 40): for param_group in optimizer.param_groups: param_group['lr'] *= 0.1 for lr_imgs, hr_imgs in train_loader: lr_imgs = lr_imgs.to(device) hr_imgs = hr_imgs.to(device) optimizer.zero_grad() pred = model(lr_imgs) # pred = lr_up + residual loss = criterion(pred, hr_imgs) loss.backward() # 梯度裁剪:防止深层网络梯度爆炸 nn.utils.clip_grad_norm_(model.parameters(), 0.4) optimizer.step() epoch_loss += loss.item() print(f"epoch {epoch+1:03d} | loss {epoch_loss / len(train_loader):.5f}")这里损失函数计算的是model输出与HR图之间的MSE。由于输出定义是lr_up + residual,该式等价于对残差项做监督;但与直接写MSE(residual, hr-lr_up)相比,前者能让你在训练过程中直接观察重建图的loss走势,更方便排查问题。
学习率衰减写死在epoch in (20, 40)处。如果你的数据规模与论文不同,比如只有几百张图,20个epoch跑一轮到收敛可能太慢,可以把衰减点按epoch比例调整到总epoch的40%、80%附近。训练量不大时通常50个epoch就能看到明显的PSNR上升趋势,我自己的经验是在第10个epoch左右loss会开始明显下降,前10个epoch的缓慢下降属于正常预热。
3.5 梯度裁剪的正确写法:PyTorch两种实现对比
PyTorch里实现梯度裁剪有两个选择,一个是nn.utils.clip_grad_norm_,它会把所有参数的梯度拼接后计算全局L2范数,然后统一缩放;另一个是逐参数裁剪,对应VDSR论文中“per-parameter gradient norm clipping”的原始描述。
# 方式一:全局梯度范数裁剪,PyTorch内建函数,最简洁 nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.4) # 方式二:逐参数裁剪,更贴近论文描述 def clip_gradient_per_param(model, threshold=0.4): for param in model.parameters(): if param.grad is not None: norm = param.grad.detach().norm() if norm > threshold: param.grad.mul_(threshold / (norm + 1e-12))两种方式在多数场景下结果区别不大,但有一个差异值得注意:全局裁剪把所有参数看成一个整体,某一个参数梯度过大时会被其他参数“平均稀释”;逐参数裁剪则保证每个参数的梯度范数都不会超过0.4,对VDSR这样所有层宽度一致的网络更贴近论文描述。我一般习惯用逐参数版本,代价是每个step多一次参数遍历,对66万参数量几乎无感。
另外要注意,梯度裁剪必须在loss.backward()之后、optimizer.step()之前执行,顺序不能颠倒。很多刚上手的人把裁剪放到step之后,结果裁剪的只是上一步的梯度,起不到任何作用。至于为什么不推荐Adam,原因是Adam自带自适应学习率,会把梯度裁剪的阈值作用复杂化,VDSR整套训练设置本来就是围绕SGD设计的,换优化器等于同时改多个超参,不利于定位问题。
4. 复现VDSR的5条避坑记录:NaN、PSNR不涨、显存不足等
4.1 梯度爆炸导致loss变成NaN:缺少梯度裁剪
现象:训练到某一轮时loss突然从正常数值跳变到NaN,之后再也恢复不回来,换更小的学习率重训也一样。
原因:二十层网络的反向传播,梯度会随层数累积放大。初始lr=1e-3在VDSR训练体系下不算小,一旦某一步输入patch正好是强纹理边缘密集的图像,梯度范数就可能超过数十,直接把权重参数顶到数值溢出。这是复现时遇到的第一个坑,当时没加梯度裁剪只调小了学习率,结果loss跌得极慢,反而拖慢了收敛。
解决:在每轮参数更新前执行梯度裁剪,阈值0.4。如果仍然NaN,把裁剪阈值降到0.1,同时检查数据加载器里是否有数值异常的图像,比如全黑或全白patch。NaN还有一个容易被忽略的来源是patch尺寸与scale不匹配导致cv2.resize拿到空矩阵,调试时在数据加载器里打印几个patch的mean和std能快速定位。想观察梯度是否逼近爆掉,可以在backward后临时打印一下梯度范数:
loss.backward() grad_norm = 0.0 for p in model.parameters(): if p.grad is not None: grad_norm += p.grad.detach().norm() ** 2 print(f"grad norm: {grad_norm ** 0.5:.4f}")正常训练时这个值会稳定在零点几到个位数,如果某一步突然跳到上百,梯度裁剪就是用来按住它的。
4.2 PSNR一直在低水平打转:残差连接没接回去
现象:训练loss一路下降,看着很顺利,但拿到测试集一算PSNR,一直在30dB左右徘徊,比论文低好几个dB。
原因:前向传播输出的是residual而不是最终重建图。如果训练时loss监督的是残差项,模型输出直接与hr-lr_up比较,loss也能下降;但到了测试阶段,如果直接拿模型输出跟HR图比PSNR,自然会差很远,因为输入的低频信息根本没被加回来。
解决:排查测试脚本,确认测试推理时使用的是model(x) + x,而不是model(x)。更稳妥的办法是把残差连接写进模型forward函数里,让所有使用模型的地方自动得到重建结果。上面3.3节的代码已经采用这种写法,可以避免这类混淆。这条坑之所以常见,是因为很多博客复现时把“训练时求残差损失”和“模型输出就是残差”混为一谈,前者没有问题,后者必须再加上输入才得到最终结果。
4.3 显存不足OOM:patch大小、batch size与多尺度采样的平衡
现象:训练刚开始报RuntimeError: CUDA out of memory,有时运行几个epoch后才出现,日志里能看到batch size大约8~16。
原因:多尺度训练时,一个batch内patch尺寸统一为64×64,数据加载器返回的lr_up和hr_patch始终是64×64,显存消耗是稳定的。真正导致OOM的通常是DataLoader的num_workers开得过大,或者GPU显存只有4GB左右,64×64×batch16在64通道的网络下已经接近上限。
解决:显存小的机器把patch_size降到48或41,batch_size降到8,必要时关闭pin_memory。如果还想保持原有patch尺寸,可以用梯度累积代替增大batch size:
accumulation_steps = 2 for step, (lr, hr) in enumerate(train_loader): loss = criterion(model(lr), hr) / accumulation_steps loss.backward() if (step + 1) % accumulation_steps == 0: nn.utils.clip_grad_norm_(model.parameters(), 0.4) optimizer.step() optimizer.zero_grad()把一次大batch拆成多次小batch累积梯度,是保留大等效batch又不爆显存的标准做法。唯一的坑是裁剪和step必须放在累积完成之后,否则等价batch的效果就失效了。另外,在WSL2里跑训练时,默认显存分配可能被Windows的桌面程序占掉一部分,优先检查一下宿主机GPU占用情况。
4.4 单尺度训练后换尺度PSNR大幅下降:数据集没做多尺度采样
现象:训练时只用scale=2,测试2x效果不错,但一测4x就明显模糊,PSNR掉到27dB以下,和论文里“一个模型应该能处理任意尺度”的说法对不上。
原因:单尺度训练让网络只在2x这一种退化模式上拟合,4x的输入纹理特征完全不同,插值放大后的LR_up与真实4x LR_up有系统性差异,网络没见过这种分布,自然表现糟糕。
解决:把数据加载器的scales参数改成(2,3,4),训练时随机采样scale。多尺度训练会让每个scale上的收敛速度稍慢,但最终得到的模型在三个尺度上都能保持稳定水平。如果你只部署在固定的2x场景,单尺度训练确实是更省算力的选择,但论文复现目标是任意尺度,这一步不能省。另一个相关经验是:多尺度训练后的模型在2x上的PSNR通常比单尺度2x略低0.1~0.2dB,这是正常现象,换取的是整体能力。
4.5 测试PSNR和论文差0.3dB以上:边界处理、像素值域、PSNR算法三处埋雷
现象:同一测试集、同一scale,自己训练的模型PSNR与论文报的数字总是差0.3~0.8dB,且怎么调学习率都补不回来。
原因:这类系统性差距通常不在模型本身,而在测试协议不统一。三处最常见的埋雷点如下。
第一是边界处理。VDSR推理时需要对整图逐像素卷积,但神经网络对图像边缘的预测质量天然低于中央区域。论文评测时会对重建图先裁掉边缘像素,常见做法是裁6像素,如果直接把整幅图像算进PSNR,边缘区域的低质量预测会拉低整体数值。第二是像素值域。如果训练时归一化到0~1,测试时用0~255的像素值去算MSE,PSNR公式里的峰值也要跟着改,混用会导致数值整体偏离。第三是PSNR公式本身,必须统一用10*log10(MAX^2/MSE),MAX为1(归一化后)或255(原始值域)。
解决:写一个统一评测函数,固定边界裁剪大小、固定值域、固定公式,然后用这个函数去评测所有对比模型,保证口径一致。第5章会给完整代码。
5. 验证训练的VDSR模型:PSNR计算、测试集对齐与模型保存
5.1 测试集选择:Set5、Set14与BSD100的特点差异
VDSR论文报告结果时最常用的三个基准集是Set5、Set14和BSD100。Set5只有5张图,以婴儿、蝴蝶、鸟类等自然图像为主,图像尺寸不大,是最快验证训练是否跑通的工具集;Set14包含14张图,加入了一些经典测试图如baboon、pepper、lena,边缘和纹理类型更丰富;BSD100的100张图以自然风景和纹理密集内容为主,PSNR数值通常比Set5低1~2dB,更接近真实场景的难度。
实操建议是先用Set5做快速迭代,训练50个epoch后去测一下PSNR是否在合理范围内,确认没问题后再用Set14和BSD100做最终验收。如果一上来就用BSD100验证,单张图尺寸大、计算时间长,调试效率很低。这三个数据集在复现超分论文时几乎是行业标准,后续跑其他SR模型也能复用同一套评测流程。
5.2 一个即插即用的PSNR测试脚本(包含边界裁剪)
测试脚本的核心是把评测口径统一,减少外部变量干扰。下面的函数接收模型、图像路径、scale和一个可选的边界裁剪宽度,输出该图的PSNR值。
import cv2 import numpy as np import torch def evaluate_psnr(model, image_path, scale, border=6, device='cuda'): """ 计算单张图像在指定scale下的PSNR。 border表示评测时裁掉图像边缘的像素数,默认6,与论文习惯一致。 """ model.eval() # 读取图像并归一化 img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) img = img.astype(np.float32) / 255.0 h, w = img.shape # 用双三次插值构造LR,再放大回原尺寸 lr = cv2.resize(img, (w // scale, h // scale), interpolation=cv2.INTER_CUBIC) lr_up = cv2.resize(lr, (w, h), interpolation=cv2.INTER_CUBIC) # 推理 with torch.no_grad(): input_tensor = torch.from_numpy(lr_up.copy()).unsqueeze(0).unsqueeze(0).to(device) pred = model(input_tensor).squeeze().cpu().numpy() # 裁剪边界后计算PSNR if border > 0: img = img[border:-border, border:-border] pred = pred[border:-border, border:-border] mse = np.mean((pred - img) ** 2) if mse < 1e-10: psnr = float('inf') else: psnr = 10 * np.log10(1.0 / mse) # 像素值域为0~1 return psnr这个脚本有两个参数说明。border的默认值6来自论文评测习惯,适用于尺度2~4的常见测试图;如果测试图本身很小,裁掉边界像素后有效面积会显著减少,这时可以降为4。psnr计算用的是10log10(1/MSE),对应归一化到0~1的值域;如果想把值域还原到0~255,公式要改成10log10(255^2/MSE),两种写法算出来的结果一致,但不要混用。
批量评估整个测试集时,用一个循环把所有图片的PSNR平均即可:
import glob for scale in [2, 3, 4]: set5_paths = sorted(glob.glob('data/test/Set5/*.png')) psnr_list = [evaluate_psnr(model, p, scale) for p in set5_paths] avg_psnr = np.mean(psnr_list) print(f"Set5 x{scale}: {avg_psnr:.2f} dB")5.3 结果对比:VDSR论文的公开数值与我们应该达到的参考区间
| 数据集 | scale=2x | scale=3x | scale=4x |
|---|---|---|---|
| Set5 | 37.53 dB左右 | 33.66 dB左右 | 31.35 dB左右 |
| Set14 | 33.03 dB左右 | 29.77 dB左右 | 28.01 dB左右 |
| BSD100 | 31.90 dB左右 | 28.82 dB左右 | 27.29 dB左右 |
上面的参考区间是VDSR论文在公开数据集上的报告值,复现时不要求精确到小数点后两位,误差控制在0.1dB以内就算合格。这里要提醒一句:训练数据规模直接影响最终PSNR。论文用了T91、General-100、DIV2K等多套数据联合训练,如果你只用T91这一套,Set5上2x测到36.8dB左右就已经说明模型学对了,不必死磕论文数值。
我自己复现时的经验是,先用T91训50个epoch,Set5 2x能达到36.5~37.2dB,再接入DIV2K扩充数据后逼近37.4dB以上。这个增量说明模型本身没问题,剩下的差距是数据量造成的。判断模型是否训练到位,可以同时看三件事:训练loss是否持续下降到平台期;测试PSNR是否接近参考区间;从测试集里挑一张图肉眼看边缘是否锐利。三者都通过,才能说复现基本成功。
5.4 模型保存与继续训练:checkpoint写入和加载
训练结束后保存模型权重和优化器状态,方便中断恢复或换数据继续训练。建议保存state_dict而不是整个模型对象,这样后续加载时不依赖Python类实例的序列化。
import torch # 保存:包含模型权重、优化器状态、当前epoch checkpoint = { 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'best_psnr': best_psnr } torch.save(checkpoint, f'vdsr_epoch{epoch+1}.pth') # 加载 checkpoint = torch.load('vdsr_epoch50.pth', map_location='cpu') model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict'])加载后如果继续训练,可以依据checkpoint['epoch']设置start_epoch,并同步学习率衰减状态;如果只做推理,只需要model_state_dict,不必加载优化器,避免内存中驻留不必要的状态变量。还有一个小习惯:加载完要把模型切到eval模式,并在推理代码里包上torch.no_grad(),否则测试时会额外计算梯度图,显存和速度都会受影响。
6. 复现之后:把VDSR封装成一个可用的超分工具
复现结束不等于终点,把它接到自己的图像处理流程里才是更实在的事。我通常会把模型包装成一个小类,输入任意尺寸的灰度图,输出放大后的HR图,内部的预处理、前向、后处理都用同一份代码,避免评测时在数据管线上再次埋雷。
class SuperResTool: """把训练好的VDSR封装成可直接调用的超分工具""" def __init__(self, checkpoint_path, device='cuda'): self.device = device self.model = VDSR() ckpt = torch.load(checkpoint_path, map_location=device) self.model.load_state_dict(ckpt['model_state_dict']) self.model.to(device).eval() def upscale(self, img_gray, scale=2): """img_gray: 0~255范围的uint8灰度图,返回同范围HR图""" h, w = img_gray.shape target_h, target_w = int(h * scale), int(w * scale) # 非整数倍缩放时先对目标尺寸取整,避免reshape错位 lr = cv2.resize(img_gray, (max(1, target_w // scale), max(1, target_h // scale)), interpolation=cv2.INTER_CUBIC) lr_up = cv2.resize(lr, (target_w, target_h), interpolation=cv2.INTER_CUBIC) # 转tensor、预测、转回numpy x = torch.from_numpy(lr_up.astype(np.float32) / 255.0).unsqueeze(0).unsqueeze(0) with torch.no_grad(): pred = self.model(x.to(self.device)).squeeze().cpu().numpy() return np.clip(pred * 255.0, 0, 255).astype(np.uint8)处理彩色图像时,不必要让网络去学RGB三通道。常见做法是把图像转到YCbCr色彩空间,只对Y通道做超分,Cb和Cr通道直接双三次放大,再合并转回RGB。原因是超分网络对亮度通道的细节重建效果最好,而色度通道人眼敏感度低,这样既省计算量也符合VDSR单通道模型的设计。
最后,如果手头没有标准测试集,我判断模型有没有白跑的小技巧是:随便拿一张包含文字或细纹理的照片裁成两半,一半原图、一半先缩小再放大,用上面这个工具跑一次超分,直接看文字的边缘是否锐利、有没有振铃。数值上PSNR可能区分不出两个模型的差异,但眼睛对文字边缘非常敏感,边缘若能看清,说明网络的残差学习已经吃到高频细节了;若边缘平滑得像水彩,八成是训练数据没吃够或者学习率衰减过早。
这篇复现笔记是几轮踩坑换来的,希望帮到你。
本文还有配套的精品资源,点击获取