简介:基于PyTorch的图像质量评估模型RankIQA源码与说明,面向计算机视觉方向需要完成课程设计或期末大作业的学生,帮助快速掌握图像质量评估任务的模型实现、训练与评估流程,并可作为高分项目直接提交。压缩包共58个文件,以40个Python脚本为核心,涵盖数据生成、模型定义、损失计算、评估指标等完整逻辑;另含6个shell运行脚本、5张样例图片、5份Markdown说明文档及依赖清单,整体大小仅257KB,结构紧凑便于查阅。源码按models、criterions、dataloader、applications等模块清晰组织,内置mobilenetv3、efficientnet等多种网络,支持排序损失与回归损失等训练方式,并配套训练、评估、JIT转换等一键运行脚本,确保下载即用无需修改。已有397人学习下载,适合想要省去环境调试与代码改动、直接获取完整可运行方案的学习者。
1. 图像质量评估与 RankIQA:排序为什么比打分更适合 PyTorch 实战
基于 PyTorch 的图像质量评估模型 RankIQA 源码,解决的是一个工程里很具体的问题:给一张图像输出一个稳定、可复现的质量分。图像质量评估(IQA)在视频编码、监控系统、图像生成管线里都是刚需,但难点在于——机器很难直接学「绝对分数」。RankIQA 的核心思路很反直觉:先别学打分,先学排序,让模型判断「图 A 比图 B 差」这种相对关系。这个 trick 在复现时比直接回归 MOS 分数好用得多,尤其适合两类人:一类是做图像处理相关课程设计、需要高质量基线模型的在校生,另一类是给图像/视频链路加质量监控模块的工程师。源码附带的课程设计说明把训练和评估流程都写清楚了,拿来当毕业论文基线和工程预研都够用。
2. RankIQA 的两阶段训练管线:合成失真、Siamese 网络与 Ranking Loss 怎么配合
2.1 为什么先学「排队」而不是直接学打分
主观 MOS 分数的标注成本很高,一张图通常要几十个人打分再取平均,而且不同人的评分习惯差异非常大——同一张图有人给 6 分有人给 8 分,这种噪声会让回归模型学得很痛苦。「这张图比那张图差」这种相对排序标注则要稳定得多,标注者的一致性可以做到 90% 以上。RankIQA 利用了这个特点:先在合成失真数据上训练排序模型,让网络把「质量感知能力」学出来,再用少量真实 MOS 数据微调成打分模型。
这个两阶段设计和迁移学习的思路是一脉相承的。排序阶段不依赖任何真实 MOS 标注,只需要合成失真图像和已知的失真等级,数据可以无限生成;微调阶段虽然需要真实标注,但量可以很小,因为在排序阶段网络已经学会了「什么样的图像特征代表质量下降」,微调只是在学「这些特征如何映射到绝对分数」。
2.2 合成失真数据集怎么生成
RankIQA 原文用的是 LIVE 数据集作为参考图来源,实际复现时,没有 LIVE 也可以用 BSDS、COCO 或者自己收集的高清自然图像代替。关键是每一张参考图都要生成多种失真类型、每个类型多个递增等级,这样任意抽两幅图就能构成一对有明确质量先后关系的训练样本。下面是我的生成脚本:
import cv2 import numpy as np import os # 4 种失真类型 × 5 个等级,等级越大失真越严重 DISTORT_TYPES = ["blur", "noise", "jpeg", "sp"] LEVELS = [1, 2, 3, 4, 5] def add_salt_pepper(img, ratio): out = img.copy() h, w = img.shape[:2] count = int(h * w * ratio) for _ in range(count): y, x = np.random.randint(0, h), np.random.randint(0, w) out[y, x] = 255 if np.random.rand() < 0.5 else 0 return out def distort(img, d_type, level): if d_type == "blur": k = 3 + level * 2 # 核大小从 5 开始,每级 +2 return cv2.GaussianBlur(img, (k, k), 0) if d_type == "noise": sigma = 5 + level * 10 # 噪声方差逐级加大 noise = np.random.normal(0, sigma, img.shape) return np.clip(img + noise, 0, 255).astype(np.uint8) if d_type == "jpeg": q = max(85 - level * 15, 10) # JPEG 质量从 70 起逐级减 15 _, enc = cv2.imencode(".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), q]) return cv2.imdecode(enc, 1) if d_type == "sp": return add_salt_pepper(img, 0.005 + level * 0.01) raise ValueError(f"unknown type: {d_type}") # 逐张读入参考图,输出为 PNG 防止二次有损压缩 for ref_name in os.listdir("ref_images"): img = cv2.imread(os.path.join("ref_images", ref_name)) if img is None: continue for t in DISTORT_TYPES: for lv in LEVELS: out = distort(img, t, lv) cv2.imwrite(f"distorted/{t}/{ref_name[:-4]}_{t}_{lv}.png", out)失真参数按「五级递进」的经验值设置:高斯模糊的核从 5 开始每级 +2,视觉差异明显但不至于完全糊掉;噪声方差从 15 起步,超过 50 时图像细节基本被噪声淹没,正好作为最高等级。这里的关键不是数值本身,而是同类型失真内部等级的视觉差异要单调——level 1 必须明显好于 level 2,否则生成的排序标签就是错的。另外注意保存格式一定要用 PNG 或 BMP,如果 JPEG 分支的结果再被保存为 JPEG 一次,相当于叠加了一种未登记的失真,会污染整个等级体系。生成完建议随机抽几组图,肉眼确认单调关系成立再进训练,这一步能省掉后面大量排查时间。
2.3 Siamese 网络与 Ranking Loss 实现
排序阶段的网络结构不复杂:一个共享权重的卷积主干(常用 VGG16),后面接一个全连接打分头,输出一个标量分数。PyTorch 里实现 Siamese 网络有个很省事的特性——不需要显式定义双塔结构,同一个模型实例跑两次 forward 就是权重共享:
import torch import torch.nn as nn import torchvision.models as models class RankNet(nn.Module): def __init__(self, backbone="vgg16", pretrained=True): super().__init__() base = models.vgg16(pretrained=pretrained).features self.features = base self.head = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(512, 128), nn.ReLU(inplace=True), nn.Linear(128, 1), # 输出单维质量分 ) def forward(self, x): return self.head(self.features(x))训练时,把同一参考图的两个不同失真版本分别送入同一个模型,得到两个分数,然后用 MarginRankingLoss 拉大它们之间的距离:
import torch.nn.functional as F def train_ranking_step(model, opt, good, bad): opt.zero_grad() s_good = model(good) # 高质量图分支输出 s_bad = model(bad) # 低质量图分支输出 target = torch.ones_like(s_good) # target=1 表示输入1应该得分更高 loss = F.margin_ranking_loss(s_good, s_bad, target, margin=1.0) loss.backward() opt.step() return loss.item()margin_ranking_loss 的公式是max(0, -target * (x1 - x2) + margin),当 target=1 时,它要求x1至少比x2高 margin 分,否则产生损失。margin 设 1.0 是常用值——margin 越大,模型越要把两个输入的分数拉开,排序越稳定,但收敛也越慢;在小数据集上可以降到 0.5 试试。注意打分头里不要加 Dropout,因为 Siamese 模式下同一个 batch 会做两次 forward,Dropout mask 不一致会让排序判断不稳定。
2.4 第二阶段:从排序模型迁移到 MOS 打分
排序模型训练完后,输出的分数只有相对意义——0.8 和 1.2 只代表后者「质量感知上更优」,不代表绝对质量分。想要得到能直接用的打分模型,需要把 stage1 的权重搬到一个单分支回归网络上,换掉打分头,用真实 MOS 数据微调。这里有一个很容易踩的坑:stage1 训练完直接拿 head 输出当 MOS 用,分数会完全没有绝对意义。
import torch import torch.nn as nn import torchvision.models as models class RankIQAReg(nn.Module): def __init__(self, pretrained_path=None): super().__init__() self.features = models.vgg16(pretrained=False).features self.reg = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(512, 256), nn.ReLU(inplace=True), nn.Dropout(0.2), nn.Linear(256, 1), # 回归头输出 MOS 预测 ) if pretrained_path: state = torch.load(pretrained_path, map_location="cpu") # 只加载 features 层权重,reg 头重新随机初始化 self.features.load_state_dict(state["features"]) def forward(self, x): return torch.squeeze(self.reg(self.features(x)), -1)加载权重时只加载 features 部分,回归头重新初始化,这个细节很关键——stage1 的 head 学的是相对排序的空间,直接拿来当回归头只会输出一堆无意义的相对值。微调阶段的常见配置是:优化器用 Adam,学习率比 stage1 低一个量级(stage1 常用 1e-4,stage2 降到 1e-5),损失函数用 SmoothL1Loss 代替 MSE,大概 10 个 epoch 左右就能稳定收敛。如果微调数据非常少(几十张),可以考虑冻结 features 前几层,只训练后面几层和回归头,防止过拟合。
3. 把 RankIQA 源码跑通:PyTorch 环境、数据准备与训练推理落地
3.1 环境与依赖版本:PyTorch 和 CUDA 版本对应关系别踩坑
整个项目跑起来不挑硬件,GTX 1060 6G 就能跑 stage1(batch 设 8),8G 以上显存可以放开到 batch 16。环境搭建最省事的路径是 conda 建一个干净环境再装 PyTorch。网上 PyTorch 安装教程很多,但真正决定能不能跑起来的不是安装命令,而是 PyTorch、torchvision、CUDA 三者之间的版本对应关系——每年都有人栽在这上面,装完之后torch.cuda.is_available()返回 False,然后开始怀疑显卡驱动。
conda create -n rankiqa python=3.9 -y conda activate rankiqa conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch -c nvidia pip install opencv-python scipy pillow tqdm| 组件 | 建议版本 | 说明 |
|---|---|---|
| Python | 3.9 或 3.10 | 3.11 也能用,但部分 CUDA 扩展编译容易出问题 |
| PyTorch | 2.x | 安装时用 conda 同时指定 pytorch-cuda,不要分开装 |
| torchvision | 与 PyTorch 严格同版本 | conda 一条命令同时装两个,天然避开版本错配 |
| opencv-python | 最新稳定版即可 | 负责图像 I/O 与失真生成 |
| scipy | 1.10+ | 计算 SROCC / LCC 评估指标 |
提示:安装完成后先跑
python -c "import torch; print(torch.cuda.is_available())"。如果输出 False,先查nvidia-smi确认驱动版本支持当前 CUDA,再查 PyTorch 是否是 CPU 版,最后才去看代码。
3.2 数据准备:目录结构、配对文件与 Dataset 加载器
数据目录的组织方式会影响后面所有脚本的写法。我的习惯是把参考图、失真图、配对清单分开存放,所有路径写入 CSV,训练时通过 Dataset 读 CSV,这样后续增加数据、切换训练集都不需要改代码:
data/ ├── ref_images/ # 原始高清参考图 ├── distorted/ │ ├── blur/ # 高斯模糊失真 │ ├── noise/ # 高斯噪声失真 │ ├── jpeg/ # JPEG 压缩失真 │ └── sp/ # 椒盐噪声失真 ├── train_pairs.csv # 排序训练对清单 └── train_mos.csv # 微调阶段 MOS 标签配对清单的生成脚本如下,文件名约定为{ref_id}_{type}_{level}.png,从文件名里解析 ref_id 和等级:
from pathlib import Path import pandas as pd import random distorted_root = Path("distorted") ref_ids = [p.name.split("_")[0] for p in distorted_root.rglob("*.png")] def list_versions(ref_id): return sorted(distorted_root.rglob(f"{ref_id}_*.png")) def level_of(path): # 文件名形如 ref0001_blur_2.png,最后一个数字就是等级 return int(path.stem.rsplit("_", 1)[1]) rows = [] for ref in set(ref_ids): versions = list_versions(ref) if len(versions) < 2: continue for _ in range(len(versions) * 5): # 每张参考图生成的配对数量 a, b = random.sample(versions, 2) if level_of(a) < level_of(b): rows.append([str(a), str(b)]) else: rows.append([str(b), str(a)]) pd.DataFrame(rows, columns=["good_path", "bad_path"]).to_csv( "train_pairs.csv", index=False)配对数量设置为len(versions) * 5,也就是每个失真版本平均参与 5 次配对,太少模型见过的相对关系不够,太多会产生大量冗余样本拖慢训练。CSV 里的 every 行都要保证 good_path 的等级严格小于 bad_path 的等级——这个约束在生成时已经通过排序保证。加载器只需要按行读取并解码图像:
from torch.utils.data import Dataset from PIL import Image import pandas as pd class RankPairDataset(Dataset): def __init__(self, pair_csv, transform): self.pairs = pd.read_csv(pair_csv) self.transform = transform def __len__(self): return len(self.pairs) def __getitem__(self, idx): row = self.pairs.iloc[idx] good = self.transform(Image.open(row.good_path).convert("RGB")) bad = self.transform(Image.open(row.bad_path).convert("RGB")) return good, bad3.3 训练启动:参数设置、日志观察与模型保存
stage1 训练脚本的骨架如下。transform 里用了 Resize + RandomCrop 的组合,相当于给模型一点随机裁剪的扰动,对排序阶段的泛化有帮助,推理时则改用中心裁剪:
import torch from torch.utils.data import DataLoader from torchvision import transforms transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) dataset = RankPairDataset("train_pairs.csv", transform) loader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=4) model = RankNet().cuda() opt = torch.optim.Adam(model.parameters(), lr=1e-4) for epoch in range(30): for step, (good, bad) in enumerate(loader): good, bad = good.cuda(), bad.cuda() loss = train_ranking_step(model, opt, good, bad) if step % 50 == 0: print(f"epoch={epoch} step={step} loss={loss:.4f}") torch.save({ "features": model.features.state_dict(), "head": model.head.state_dict(), }, f"stage1_epoch{epoch}.pth")- batch_size:16 对应约 8G 显存;显存不够就降到 8,别改模型结构。
- 学习率:Adam 默认 1e-4 起步,stage1 不建议超过 5e-4,否则 loss 会震荡甚至发散。
- epoch:30 是个够用的值,实际看 loss 曲线,连续 3 个 epoch 不下降就可以早停。
- 权重以 dict 形式分开保存 features 和 head,是为了 stage2 微调时只加载 features,head 重新初始化。
训练日志里 loss 不需要降到 0——margin ranking loss 降到 0.3 左右、且模型对同一批图输出的分数范围有明显方差,就是健康信号。如果 loss 降得很快但分数输出几乎不变,说明模型在偷懒,需要回头调整结构。
3.4 推理:对单张图像输出质量分数
stage2 微调完成后,推理脚本非常短。核心是预处理必须和训练时保持一致——同一套 resize 尺寸、同一套 Normalize 参数:
from PIL import Image import torchvision.transforms as T def predict_score(model, img_path, device="cuda"): model.eval() x = T.Compose([ T.Resize((224, 224)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])(Image.open(img_path).convert("RGB")).unsqueeze(0).to(device) with torch.no_grad(): return model(x).item() # score = predict_score(reg_model, "test/example.png") # score 越大表示质量越好;多张图比较时必须保持同一输入尺寸打分模型的输出没有固定的零点,绝对值意义有限。实际工程里最可靠的用法是拿它做「同一批图的质量排序」——比如视频抽帧后筛选最清晰的帧,或者在生成模型的多张候选输出里挑质量最高的那张。跨批次的分数比较要谨慎,除非你做了分数校准。
4. RankIQA 避坑指南:五个训练翻车现场与排查思路
排序模型看着简单,跑起来全是细节。下面五个坑是按出现频率排的,每一条都来自实际训练日志,现象、原因、解决一条条说清楚。
4.1 loss 降了但输出没有区分度
现象:margin ranking loss 一路降到 0.01,但把测试图输入模型,所有 score 输出都是同一个常数。
原因:margin loss 为 0 只代表「在这个 batch 里没有违反排序」,不代表模型学到了区分能力。最常见的情况是学习率过大,输出直接饱和到一个值;或者打分头的隐藏单元太少,容量不够。
解决:训练中定期打印一个 batch 的 score 标准差。如果标准差趋近于 0,先把学习率降到 1e-5 重新训练,同时把打分头从Linear(512, 128)+Linear(128, 1)改成宽一点的结构。我一般用512 -> 256 -> 1,并且保证 head 里没有 Dropout。
4.2 stage2 微调后 SROCC 反而下降
现象:stage1 排序训练的分数很有区分度,但用 MOS 微调后,SROCC 比 stage1 还低。
原因:MOS 标注的尺度(比如 1-5 分或 0-100 分)和模型输出范围不匹配。回归头直接去拟合绝对分数,梯度的尺度被 label 放大,训练不稳定。另一个常见原因是 MOS 标注本身噪声大,MSE loss 对异常标注非常敏感。
解决:训练前对 MOS 做 z-score 标准化(减均值除方差),预测后逆变换还原分数。损失函数换成 SmoothL1Loss,它对离群点比 MSE 更鲁棒。微调学习率控制在 1e-5 量级,别超过。
4.3 数据随机划分导致验证集结果虚高
现象:train/val 按文件随机划分时,验证 SROCC 能到 0.97,换到一个全新的数据集上直接崩到 0.7。
原因:同一参考图的多个失真版本来自同一样本,内容高度相关。随机划分会把同一参考图的失真版本同时分到训练集和验证集,模型相当于「见过」验证集的内容,评估结果虚高。
解决:按参考图分组划分数据——同一个 ref_id 的所有失真版本,要么全部进 train,要么全部进 val。这是 IQA 数据划分的硬规则,没有例外。写划分脚本时用 ref_id 作为分组键,不要用单个文件路径。
4.4 JPEG 失真生成库不一致导致等级错乱
现象:在同一台机器上用 OpenCV 生成 JPEG 失真训练,到另一台机器用 PIL 复现,发现相同 quality 参数的图像视觉质量明显不同,等级顺序对不上。
原因:OpenCV 和 PIL 底层用的 libjpeg 版本和量化表不一样,相同 quality 参数不代表相同视觉质量,出来的压缩痕迹完全不同。
解决:整个项目固定只用一种编码库生成训练数据和测试数据,我习惯统一用 OpenCV。生成之后抽几张图计算 PSNR,确认每个失真类型内部等级严格单调——level 越高 PSNR 越低,这是最直接的自检手段。
4.5 显存不够:VGG16 双输入是显存杀手
现象:batch_size 设 16,训练刚开始就报 CUDA out of memory。
原因:Siamese 网络虽然共享权重,但两次 forward 的中间激活都要存在显存里,实际显存占用约等于两个 VGG16 同时前向,VGG16 本身又是出了名的显存黑洞。
解决:降 batch 到 8,输入裁剪到 224×224,这通常就够了。如果还爆,用torch.utils.checkpoint包住 features 层,以少量计算换显存。想快速验证流程正确性,先用 mobilenet_v2 跑通再换 VGG16,能省很多调试时间。
5. 进阶用法:多尺度特征回归与跨数据集评估的验证技巧
5.1 把最后一层换成中间层拼接特征
RankIQA 原版在 VGG16 features 后接单个打分头,实践中我发现换一种接法对跨数据集泛化有稳定提升:取网络浅层和深层的池化特征拼接后再回归。浅层特征保留纹理和边缘信息,深层特征捕捉语义退化,两者互补:
import torch import torch.nn as nn import torchvision.models as models class MultiFeatRankIQA(nn.Module): def __init__(self): super().__init__() base = models.vgg16(pretrained=False).features self.low = base[:17] # 到 conv3_3,输出 256 通道 self.high = base[17:] # conv4 至 conv5_3,输出 512 通道 self.pool = nn.AdaptiveAvgPool2d(1) self.reg = nn.Sequential( nn.Linear(256 + 512, 128), nn.ReLU(inplace=True), nn.Linear(128, 1), ) def forward(self, x): low_feat = self.low(x) high_feat = self.high(low_feat) f_low = self.pool(low_feat).flatten(1) f_high = self.pool(high_feat).flatten(1) return self.reg(torch.cat([f_low, f_high], dim=1)).squeeze(-1)这里的维度对应 VGG16 的固定结构:base[:17]正好切到 conv3_3 输出 256 通道,base[17:]从 conv4_1 开始到 conv5_3 输出 512 通道。如果你换 backbone,务必先打印每层输出形状再改拼接维度,这个维度错误很难一眼发现。训练配置和单特征版完全一致,直接替换模型类即可。
5.2 SROCC / LCC / RMSE:三个指标怎么配合看
RankIQA 论文和大多数 IQA 工作都以 SROCC 作为主指标,但工程落地时只看一个指标会误判:
| 指标 | 全称 | 关心什么 | 什么时候看它 |
|---|---|---|---|
| SROCC | Spearman 秩相关系数 | 预测与 MOS 的单调一致性 | 无参考场景、混合失真类型时 |
| LCC | Pearson 线性相关系数 | 预测与 MOS 的线性相关程度 | 需要绝对分数校准、同一数据集内比较 |
| RMSE | 均方根误差 | 预测与真实分数的偏差大小 | 做分数阈值判断、质量门控时 |
计算方式用 scipy 一行搞定,预测分数和 MOS 列表对齐即可:
from scipy.stats import spearmanr, pearsonr srcc, _ = spearmanr(pred_scores, mos_labels) lcc, _ = pearsonr(pred_scores, mos_labels) rmse = ((pred_scores - mos_labels) ** 2).mean() ** 0.55.3 跨数据集评估的现实边界
在合成失真上训练、在真实相机失真数据上评估时会有明显的分数漂移,这是合成失真与真实失真分布差异导致的,不是代码 bug。我一般会在评估报告里同时给出「同源数据集 SROCC」和「跨数据集 SROCC」,并标注清楚微调和未微调的区别。做真实场景落地时,收集几十张真实失真图做 stage2 微调,效果提升往往比优化网络结构更明显。
从那以后,我每次接到新的 IQA 项目,都会强制先跑一遍「合成失真 + ranking 预训练 → 小规模 MOS 微调 → SROCC 报告」这条基线流程,确认数据划分、失真生成、损失函数三个环节都没问题,再谈改进模型结构。这个 RankIQA 项目本身就是一条很扎实的基线,跑通一次,你对排序学习、迁移训练和图像质量评估的完整认知就建立起来了。希望帮到你。
本文还有配套的精品资源,点击获取