简介:面向计算机视觉课程设计与期末大作业,提供了一套基于PyTorch实现的RankIQA图像质量评估模型完整源码。项目采用排序学习机制训练无参考图像质量评估模型,包含从数据准备、模型构建、损失函数设计到训练评估的完整流程,特别适合需要高分课程设计或希望系统学习RankIQA算法的学生与开发者。资源包共包含五十八个文件,以四十个Python脚本为核心,覆盖数据生成、模型定义、损失计算、训练测试等模块;同时附有shell运行脚本、Markdown说明文档、图片示例及环境配置要求,整体压缩包仅257KB,十分轻量,便于直接迁移。目前已有三百九十七人学习或下载,说明该资源经过了不少实际使用检验,内容可靠。源码内含项目配置、自建数据加载器、多种损失函数(排序损失、回归损失、EMD损失)以及多种经典模型结构,并配有运行说明,下载即可运行,无需修改;既可直接作为高分课程设计提交,也可作为期末大作业,便于在报告中详细分析实现细节。
1. 基于PyTorch的图像质量评估模型RankIQA:先学会排序,再学会打分
做图像质量评估(IQA)的人都有一个共同的痛:想训练一个能预测主观评分的模型,需要大量人工标注的MOS分数,而标分数既贵又不稳定,不同人打分差异极大。RankIQA的思路恰恰是绕开这个难点——它先让模型在图像排序对上学“哪张图质量更好”,然后再用少量带分数的图像把排序能力校准成打分能力。这个思路在课程设计和实际项目里都很讨巧:不需要人工标几百张图的分数,只需要构造相对顺序,模型就能学会质量表征。这篇文章会按“原理 → 数据构造 → 模型搭建 → 踩坑 → 进阶校准”的顺序,把一套可以直接复现的PyTorch实现完整讲透,适合正在做课程设计、毕业设计,或者想快速在IQA方向落地一个可运行方案的人。本文基于PyTorch实现RankIQA,所有代码都可以直接跑通。
2. RankIQA的核心策略:为什么“排序”比“打分”更容易学
2.1 从分数回归到排序学习:RankIQA到底改了什么
传统的IQA模型,比如经典的BRISQUE、NIQE,或者后来基于深度网络的模型,直接学习“图像 → MOS分数”的映射。这里的问题在于:MOS分数本身是主观平均值,噪声很大,同一张图像在不同实验里得到的分数可能差出1分以上。模型拟合这种带噪声的标签,要么过拟合到个别标注者的偏好上,要么学到的表征不够泛化。
RankIQA做了一个关键转换:把回归问题变成排序问题。模型不再预测“这张图值3.2分”,而是预测“这张图比那张图质量高”。排序对的标签是稳定的——对同一场景,加噪声越重、压缩越狠的图像,质量一定更差,这个顺序几乎不会因为标注者不同而改变。用ResNet等预训练网络做特征提取,再在特征之上加一个排序头,用Pairwise Ranking Loss训练,模型就能学到对质量敏感的表征。这一步的妙处在于,排序对的构造完全不需要人工打分,只需要知道“哪张图是原始图、哪张图是失真图”。
2.2 两阶段训练:排序预训练 + 回归微调
RankIQA原文采用了两阶段方案,这个设计直接影响源码结构:
阶段一是排序学习。构建孪生网络(Siamese Network),两个分支共享权重,输入是一对图像,输出各自的预测分数,用一个排序损失函数约束“干净图的输出 > 失真图的输出”。这一步让网络学会质量排序。
阶段二是回归微调。把排序头的输出换成回归头,用少量有真实MOS标签的图像做监督微调。因为网络已经在排序任务上获得了良好的质量感知表征,微调只需要很少的标注数据就能达到不错的效果。
我在实际复现时发现,阶段二不是可选项,而是必选项。如果不做回归微调,模型只能输出相对排序结果,没法给出一张图具体的质量分数,也就没法计算PSNR之外的评价指标。课程设计要交“预测分数 vs 真实分数”的对比实验,阶段二必须做。
提示:两阶段共享同一个特征提取器,但阶段二需要换掉最后的全连接层,尺寸从“1个神经元”换回“1个神经元”在结构上没区别,但训练方式要从Pairwise Loss换成L1或MSE Loss。
2.3 网络结构怎么选:ResNet18还是ResNet50
RankIQA的特征提取器没有硬性规定,常见做法是直接加载PyTorch预训练的ResNet。选型时考虑三点:
显存与速度。孪生网络一次前传要跑两张图,ResNet50比ResNet18显存占用高出一倍多。课程设计一般用单卡,ResNet18在1080Ti上训练256×256的patch,batch_size 16勉强能跑,ResNet50就得降到8。
预训练权重的影响。ImageNet预训练权重对IQA任务有正面帮助,因为卷积核已经学会了纹理、边缘、结构等底层特征,质量退化恰恰体现在这些特征上。但要注意,预训练权重的末层分类头(1000类)必须换掉。
微调的灵活性。阶段二微调时,冻结前几层、只解冻后面几个stage的效果通常更好。ResNet的分层结构更适合这种局部解冻操作,这也是我优先选它的原因。
3. 构造训练数据:把无标签图像变成排序对
3.1 失真类型与数据集选择:LIVE、TID2013还是自己造
RankIQA的训练数据主要是“原始图像 + 各类失真图像”。常用的公开数据集有LIVE(包含JPEG压缩、高斯噪声、高斯模糊等失真)、TID2013(失真类型更丰富,含色彩失真和频段噪声)。课程设计如果只想跑通流程,用LIVE就够了;如果论文里想强调泛化性,TID2013是更好的选择。
但有个现实问题:LIVE的原始图像只有30张左右,全部是25mm胶片时代的照片,分辨率不高。直接用原始图构造排序对,样本多样性不够。常见做法是先从原始图上随机裁剪patch,再对每个patch施加不同强度的失真。这样一对图像从同一patch来,内容完全相同,只有失真程度不同,排序标签绝对可靠。
3.2 Patch切分与失真施加:最小可执行代码
我一般用OpenCV和NumPy直接写数据准备脚本,不依赖额外的IQA工具库。下面这段代码做的事情是:读取一张干净图,随机裁剪出固定大小的patch,然后对它施加两种强度的JPEG压缩失真,形成一个排序对。代码里失真类型可以替换成高斯噪声、高斯模糊,方法一样。
import cv2 import numpy as np def random_crop(img, crop_size=224): h, w = img.shape[:2] y = np.random.randint(0, max(1, h - crop_size)) x = np.random.randint(0, max(1, w - crop_size)) return img[y:y + crop_size, x:x + crop_size] def jpeg_distort(patch, quality): # quality 越小,压缩越狠,质量越差 encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality] _, enc = cv2.imencode('.jpg', patch, encode_param) return cv2.imdecode(enc, 1) def make_pair(clean_img, crop_size=224, q_high=80, q_low=20): patch = random_crop(clean_img, crop_size) # 原图patch ref = jpeg_distort(patch, q_high) # 高质量版 dist = jpeg_distort(patch, q_low) # 低质量版 return ref, dist # ref 的质量一定高于 dist逻辑说明:随机裁剪保证两张图内容对齐,JPEG重压缩会产生块效应,quality参数越小块效应越明显。排序对的顺序是确定的——ref在前、dist在后,训练时标签固定为“前者的质量分数必须大于后者”。这里务必保证输入的BGR格式一致,OpenCV读出来是BGR,PyTorch的预训练模型要求RGB,后面转换时要统一。
参数说明:crop_size通常取224或256,要和ResNet的输入尺寸匹配;q_high和q_low是质量边界,差距越大排序任务越简单,但差距过大模型容易学不到细粒度差异。实际训练中我建议把quality区间设成[20, 90],每对随机取两个值,不要固定80和20。
3.3 批量生成排序对的完整脚本
单张图最多只能生成有限的patch,要撑起一个epoch至少几千对,必须批量处理整个数据集。下面这段脚本把上述操作封装成数据集类,直接对接PyTorch的DataLoader。
import torch from torch.utils.data import Dataset import random class RankPairDataset(Dataset): def __init__(self, image_paths, crop_size=224, distort_fn=jpeg_distort, quality_range=(20, 90), transform=None): self.paths = image_paths self.crop_size = crop_size self.distort_fn = distort_fn self.quality_range = quality_range self.transform = transform # 必须做ToTensor和ImageNet归一化 def __len__(self): return len(self.paths) * 8 # 每张图生成8个patch def __getitem__(self, idx): img = cv2.imread(self.paths[idx % len(self.paths)]) patch = random_crop(img, self.crop_size) # 随机取两个质量值,保证q1 > q2 q1, q2 = random.sample(range(*self.quality_range), 2) if q1 < q2: q1, q2 = q2, q1 ref = self.distort_fn(patch, q1) dist = self.distort_fn(patch, q2) if self.transform: ref = self.transform(ref) dist = self.transform(dist) # 标签:1 表示 ref 比 dist 质量高 return ref, dist, torch.tensor(1.0)逻辑说明:__getitem__里每次随机采样两个质量参数,保证顺序正确。标签固定为1.0,配合Ranking Loss使用时表示“第一个输入的质量得分应高于第二个”。transform里必须带ImageNet的均值方差归一化,否则预训练权重直接废掉。
参数说明:len(self.paths) * 8是经验值,patch越多一个epoch越长,训练越稳,但数据加载时间也越长。如果机器性能紧张,改成*4即可。quality_range建议下限至少15,低于15的JPEG压缩图像几乎没有可用视觉信息,模型学不到有意义的结构特征。
4. 搭建RankIQA模型并训练:PyTorch实现与参数调优
4.1 孪生网络与Ranking Loss:核心代码逐行拆解
RankIQA的排序阶段网络结构很简单:共享的ResNet特征提取器 + 一个全连接输出头。PyTorch实现孪生网络的常见做法不是定义两个独立模型,而是定义单模型,输入时把两张图拼成batch一起前传,再拆开计算损失。这样省显存,也避免两个分支权重不同步的问题。
import torch import torch.nn as nn import torchvision.models as models class RankIQA_Siamese(nn.Module): def __init__(self, base='resnet18', feat_dim=512): super().__init__() if base == 'resnet18': backbone = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) elif base == 'resnet50': backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1) # 去掉原始分类头 self.features = nn.Sequential(*list(backbone.children())[:-2]) self.pool = nn.AdaptiveAvgPool2d(1) # 排序头:输出1个标量代表质量分数 self.head = nn.Sequential( nn.Linear(feat_dim, 128), nn.ReLU(inplace=True), nn.Linear(128, 1) ) def forward(self, x): x = self.features(x) x = self.pool(x) x = torch.flatten(x, 1) return self.head(x)逻辑说明:self.features保留了ResNet从conv1到layer4的全部卷积层,最后两层(AvgPool和FC)被移除,改用AdaptiveAvgPool2d(1)加自定义头。这样不管输入分辨率是多少,池化后特征图都是1×1,feat_dim只由输出通道数决定——ResNet18是512,ResNet50是2048。
参数说明:feat_dim必须和backbone的输出通道严格对应,写错会在forward时报维度不匹配。这里的排序头是一个两层的MLP,中间隐层128维是平衡容量和速度的选择;如果训练数据量大,可以放宽到256,但Phase 2微调时隐层太宽容易过拟合少量分数标签。
Ranking Loss使用MarginRankingLoss,这是PyTorch内置的排序损失,正好适用于“输入一对得分,约束第一个大于第二个”的场景。
criterion = nn.MarginRankingLoss(margin=0.1) def rank_loss_fn(ref_score, dist_score, target): # target 全是 1.0,表示 ref_score 应大于 dist_score return criterion(ref_score, dist_score, target)MarginRankingLoss的计算方式是max(0, -target * (score1 - score2) + margin)。当target=1时,损失惩罚“ref_score ≤ dist_score”的情况,margin控制两个分数的间隔要求。margin设太小(0.01)模型分不清细微质量差,设太大(1.0)训练初期loss降不下去,0.1是一个稳妥的起点。
4.2 训练循环:两阶段分别怎么跑
排序阶段的训练循环和普通分类任务区别不大,但有几个细节必须注意:两张patch前传后要拆开计算损失,梯度更新的频率要和batch_size匹配,学习率要比常规分类任务低。
def train_rank(model, dataloader, epochs=10, lr=1e-4, device='cuda'): model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.5) criterion = nn.MarginRankingLoss(margin=0.1) model.train() for epoch in range(epochs): total_loss = 0.0 for ref, dist, target in dataloader: ref, dist, target = ref.to(device), dist.to(device), target.to(device) # 拼成batch一次前传,省显存 x = torch.cat([ref, dist], dim=0) scores = model(x) ref_score, dist_score = scores.chunk(2, dim=0) loss = criterion(ref_score, dist_score, target) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() scheduler.step() print(f'Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(dataloader):.4f}')逻辑说明:torch.cat([ref, dist], dim=0)把一对图像在batch维度拼起来,前传一次,再用chunk(2)拆开。这样网络每个batch只前传一次而不是两次,计算图更简洁,反向传播也更稳定。target在这里是标量张量,全为1.0,表示ref_score要大于dist_score。
参数说明:学习率1e-4是排序阶段的常用值,比常规分类低一个数量级,因为预训练权重已经足够好,学习率太大会破坏底层特征。step_size=4, gamma=0.5的意思是每4个epoch学习率减半,训练10个epoch总共衰减两次。如果loss在后期震荡,把这个scheduler换成ReduceLROnPlateau会更灵活。
阶段二的微调代码在结构上和上面几乎一样,只有三处区别:损失函数换成L1Loss或MSELoss,数据加载换成带真实MOS分数图像,优化器学习率降到1e-5级别并冻结backbone的前几层。下面这段是微调的核心片段。
def train_regression(model, train_loader, val_loader, epochs=20, lr=1e-5): # 冻结前三个stage,只训练最后一层 + 回归头 for name, param in model.named_parameters(): if 'layer4' not in name and 'head' not in name: param.requires_grad = False optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr) criterion = nn.L1Loss() for epoch in range(epochs): for img, mos in train_loader: img, mos = img.to(device), mos.to(device) pred = model(img) loss = criterion(pred.squeeze(), mos) optimizer.zero_grad() loss.backward() optimizer.step()逻辑说明:'layer4' not in name and 'head' not in name这段判断是冻结策略的核心。训练过程中只更新layer4和最后的head,其他参数保持不变。这样做的原因是排序阶段已经让底层特征具有质量感知能力,回归阶段只需要在顶层把相对分数映射到MOS空间。全部解冻会让模型在少量标注上过拟合。
参数说明:L1Loss比MSE Loss对离群标注更鲁棒,MOS分数标注不稳定,个别离群点不应该主导梯度。lr=1e-5对只训练最后两层来说仍然偏保守,如果loss在验证集上不动,可以考虑解冻layer3,但learning rate要降到5e-6。
4.3 训练过程中的验证策略与收敛判断
RankIQA训练不能只看loss,因为排序loss下降只说明“模型能把图像分出高下”,不代表分数准。要在验证集上同时看两个指标:排序准确率(pairwise accuracy)和SROCC(Spearman秩相关系数)。
排序准确率的计算方式很直观:取一批测试图,对每个场景生成一组不同失真的图像,模型给它们打分,然后计算预测排序和真实失真强度排序的一致率。SROCC则需要真实MOS分数,适合在有MOS标注的数据集上评估。
我一般每隔一个epoch在LIVE数据集上算一次SROCC,如果SROCC连续3个epoch不涨,就把学习率降到原来的三分之一。比起训练loss,SROCC更接近最终评价指标,也更不容易被rank loss的偶然波动误导。
提示:排序阶段结束时模型是不输出具体MOS分数的,数字大小只有相对意义。不要在这个阶段把输出当真实分数去算PSNR类的误差指标,那是阶段二微调之后才能做的事。
5. RankIQA实战避坑:训练翻车最常见的5个原因
5.1 排序loss不降反升,准确率卡在50%上下
现象:训练了十几个epoch,排序loss在0.5附近震荡,验证集的pairwise accuracy始终在55%以下,和随机猜差不多。
原因:最大概率是数据构造时排序对“内容不一致”。如果裁剪patch时没固定随机种子,或者两张图不是从同一位置裁剪出来的,模型学到的是“内容不同带来的分数差异”,而不是“质量不同带来的分数差异”。内容差异远大于质量差异时,排序任务就变成了图像检索任务,模型完全跑偏。
解决:检查数据加载代码,确保ref和dist来自同一次random_crop的返回值。我在上面给的代码里是先裁剪一次再分别压缩,这就是标准做法。如果你从预处理缓存里读图,确认两张图是同一裁剪位置的产物。另一类原因是quality_range上下界差距过小,比如[70, 80],这种区间内JPEG压缩的视觉差异微乎其微,排序标签虽然正确但信号太弱,把区间扩大到[15, 90]重新训练。
5.2 验证集SROCC很高,但实际图像评分效果很差
现象:LIVE数据集上SROCC达到0.93,看起来不错,但你拿自己拍的照片去测,分数分布完全不合理,比如严重模糊的照片反而得了高分。
原因:这是典型的“数据集偏差”。LIVE的失真类型有限,主要是JPEG压缩、高斯噪声、高斯模糊和快衰落。模型在训练时只见过这几种失真,遇到真实场景中的运动模糊、过曝、低光照噪点时,特征分布超出训练分布,预训练网络会把这些当作另一种“干净图”。
解决:训练时混合多种失真类型,不要只用JPEG。在数据准备脚本里增加高斯噪声、高斯模糊、椒盐噪声几种失真函数,每个batch随机选择一种。TID2013里包含24种失真,课程设计如果时间允许,建议至少覆盖5到6种常见的。另一个手段是数据增强里加入随机缩放和亮度抖动,提高模型对内容变化的鲁棒性。
5.3 显存爆炸,batch_size调到4都跑不起来
现象:用ResNet50训练孪生网络,batch_size设16,程序启动没多久就报CUDA out of memory。
原因:孪生网络虽然只前传一次,但一次前传就吃了2倍batch的图。而且torch.cat([ref, dist])让batch维度变成2N,ResNet50在224×224输入下,2×16=32张图的前传显存占用接近8GB。如果再开Adam的动量缓存,显存直接翻倍。
解决:换ResNet18,特征维度降一半,显存占用直降。或者把batch_size降到8,同时把pin_memory=True和num_workers=4加上,用数据加载速度弥补batch大小的不足。如果还想继续用ResNet50,可以试一下梯度累积:每4个小batch累计一次梯度再更新,效果等同于batch_size扩大4倍,显存不变。
# 梯度累积写法:每 accumulation_steps 个batch更新一次 accumulation_steps = 4 optimizer.zero_grad() for i, (ref, dist, target) in enumerate(dataloader): # ... forward 和 loss 计算 ... loss = loss / accumulation_steps # 先归一化 loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()逻辑说明:loss / accumulation_steps是梯度累积的标准做法,把每个小batch的梯度缩放到等效大batch的尺度,避免梯度累计导致更新步长异常。这个写法对孪生网络特别有用,因为它天然把单batch的显存需求翻倍。
参数说明:accumulation_steps=4表示每4个小batch才更新一次权重,训练实际effective batch size是4 × batch_size,学习率可以相应调高一点,但不要超过两倍,否则收敛不稳定。
5.4 加载预训练权重时报尺寸不匹配
现象:代码里models.resnet18(weights=...)后修改了head,但保存和加载模型时load_state_dict报size mismatch for head.0.weight之类的错误。
原因:PyTorch的load_state_dict默认要求key和shape完全匹配。你把ResNet的最后一层FC换掉了,预训练权重里没有新head的参数;反过来,如果你保存了整个模型,fine-tune时加载也会遇到head尺寸不一致的问题,因为排序头和回归头可能结构不同。
解决:加载权重时加上strict=False,让缺失或不匹配的参数保持随机初始化。只保留匹配的预训练参数。
pretrained = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) model = RankIQA_Siamese(base='resnet18') # 只复制卷积层权重,跳过FC层 pretrained_dict = pretrained.state_dict() model_dict = model.state_dict() # 过滤掉形状不匹配的key pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and model_dict[k].shape == v.shape} model_dict.update(pretrained_dict) model.load_state_dict(model_dict)逻辑说明:这段代码用字典推导式把预训练权重中“key存在且shape一致”的参数筛选出来,更新进网络当前权重。这样head层不会被错误赋值,而卷积层的预训练权重能正确加载。相比strict=False,这种方式更安全——strict=False会静默忽略所有缺失key,一旦特征层名字写错也不会报错,问题会被掩盖到训练阶段才暴露。
参数说明:如果你改动了backbone中间层的结构(比如把layer4的stride改成1),shape极可能大面积不匹配,此时需要打印pretrained_dict和model_dict对比,确认哪些层没加载成功。不要跳过这一步直接训练。
5.5 预测分数集中在一个狭窄区间,区分度不够
现象:阶段二微调完成后,预测分数在3.4到3.6之间波动,而真实MOS分数范围是0到5,模型几乎把所有图都判成“中等质量”。
原因:这是L1Loss加少量标注数据的通病。当标注样本少、训练epoch多时,模型学到的是“输出训练集MOS均值”,这是一种最安全的回归策略——因为L1Loss在预测值接近中位数时总误差最小。排序阶段带来的区分能力没有传导到回归头,可能因为回归头训练时学习率太高,把排序头学到的相对分数差距抹平了。
解决:两个手段配合使用。第一,回归微调前,先把排序模型的输出做一次线性归一化,让模型输出的分数范围映射到[0, 5]区间,再用这个归一化后的分数初始化回归头的bias。第二,降低学习率到1e-6,回归头只训练5到10个epoch,提前用验证集SROCC选模型,不要傻跑完预设epoch。如果条件允许,给回归训练数据加一张“质量极差图”和一张“质量极好图”作锚点,能有效拉伸分数分布。
6. 把RankIQA用到新场景:分数校准与可复现性验证
排序模型本质上输出的是一个相对分数,它在不同数据集上的分布不同。换到新场景时,不需要重新训练整个模型,只需要在少量新场景图像上做一次分数校准。常见的校准做法是线性映射:在新数据上采样几十张图像,人工或基于算法给一个粗略排序,然后按最小二乘拟合一条y = ax + b的线性映射,把排序模型输出映射到目标评分区间。注意,这里的采样图像要覆盖从极差到极好的全分布,只取中间段会让拟合出的斜率失真。
验证模型是否值得采纳,我个人的习惯是计算SROCC和PLCC两个指标,而不是只看单张图的预测误差。SROCC衡量排序一致性,PLCC衡量线性相关性。RankIQA的典型优势是SROCC高但PLCC可能偏低,因为排序学习不保证绝对分数线性。如果课程设计的汇报里被问到“为什么PLCC没那么高”,可以直接说:PLCC低说明回归校准还不够好,可以通过增加校准样本或引入非线性映射(比如sigmoid)解决。
复现性方面,把随机种子固定是血泪经验。PyTorch的DataLoader如果不开shuffle=False,每次epoch的数据顺序都不同,排序对的质量参数也是随机采样的,这会导致训练结果难以稳定复现。需要在脚本开头固定所有随机源,包括Python、NumPy和PyTorch的CUDDA后端。我一般把固定种子的代码写成一个函数,每个训练脚本第一行就调用,包括排序阶段和回归阶段都要固定,否则微调的结果对不上就没法定位是参数问题还是随机问题。
动手实验时先跑一个小的smoke test——用10张图、每个patch只裁剪一次、训练1个epoch,确认前向和反向传播没问题,再去跑完整的数据集。这样可以避免在花了两小时构建的完整数据集上发现一个低级错误。RankIQA的实现并不复杂,但它的两阶段设计决定了出问题的地方往往在数据构造和阶段衔接上,把上面五个坑避开,这套方案是能稳定到达“课程设计高分项目”水平的。希望帮到你。
本文还有配套的精品资源,点击获取