简介基于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(funknown 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(fdistorted/{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, backbonevgg16, pretrainedTrue): super().__init__() base models.vgg16(pretrainedpretrained).features self.features base self.head nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(512, 128), nn.ReLU(inplaceTrue), 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) # target1 表示输入1应该得分更高 loss F.margin_ranking_loss(s_good, s_bad, target, margin1.0) loss.backward() opt.step() return loss.item()margin_ranking_loss 的公式是max(0, -target * (x1 - x2) margin)当 target1 时它要求x1至少比x2高 margin 分否则产生损失。margin 设 1.0 是常用值——margin 越大模型越要把两个输入的分数拉开排序越稳定但收敛也越慢在小数据集上可以降到 0.5 试试。注意打分头里不要加 Dropout因为 Siamese 模式下同一个 batch 会做两次 forwardDropout 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_pathNone): super().__init__() self.features models.vgg16(pretrainedFalse).features self.reg nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(512, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.2), nn.Linear(256, 1), # 回归头输出 MOS 预测 ) if pretrained_path: state torch.load(pretrained_path, map_locationcpu) # 只加载 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-4stage2 降到 1e-5损失函数用 SmoothL1Loss 代替 MSE大概 10 个 epoch 左右就能稳定收敛。如果微调数据非常少几十张可以考虑冻结 features 前几层只训练后面几层和回归头防止过拟合。3. 把 RankIQA 源码跑通PyTorch 环境、数据准备与训练推理落地3.1 环境与依赖版本PyTorch 和 CUDA 版本对应关系别踩坑整个项目跑起来不挑硬件GTX 1060 6G 就能跑 stage1batch 设 88G 以上显存可以放开到 batch 16。环境搭建最省事的路径是 conda 建一个干净环境再装 PyTorch。网上 PyTorch 安装教程很多但真正决定能不能跑起来的不是安装命令而是 PyTorch、torchvision、CUDA 三者之间的版本对应关系——每年都有人栽在这上面装完之后torch.cuda.is_available()返回 False然后开始怀疑显卡驱动。conda create -n rankiqa python3.9 -y conda activate rankiqa conda install pytorch torchvision pytorch-cuda11.8 -c pytorch -c nvidia pip install opencv-python scipy pillow tqdm组件建议版本说明Python3.9 或 3.103.11 也能用但部分 CUDA 扩展编译容易出问题PyTorch2.x安装时用 conda 同时指定 pytorch-cuda不要分开装torchvision与 PyTorch 严格同版本conda 一条命令同时装两个天然避开版本错配opencv-python最新稳定版即可负责图像 I/O 与失真生成scipy1.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, indexFalse)配对数量设置为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_size16, shuffleTrue, num_workers4) model RankNet().cuda() opt torch.optim.Adam(model.parameters(), lr1e-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(fepoch{epoch} step{step} loss{loss:.4f}) torch.save({ features: model.features.state_dict(), head: model.head.state_dict(), }, fstage1_epoch{epoch}.pth)batch_size16 对应约 8G 显存显存不够就降到 8别改模型结构。学习率Adam 默认 1e-4 起步stage1 不建议超过 5e-4否则 loss 会震荡甚至发散。epoch30 是个够用的值实际看 loss 曲线连续 3 个 epoch 不下降就可以早停。权重以 dict 形式分开保存 features 和 head是为了 stage2 微调时只加载 featureshead 重新初始化。训练日志里 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, devicecuda): 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(pretrainedFalse).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(inplaceTrue), 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], dim1)).squeeze(-1)这里的维度对应 VGG16 的固定结构base[:17]正好切到 conv3_3 输出 256 通道base[17:]从 conv4_1 开始到 conv5_3 输出 512 通道。如果你换 backbone务必先打印每层输出形状再改拼接维度这个维度错误很难一眼发现。训练配置和单特征版完全一致直接替换模型类即可。5.2 SROCC / LCC / RMSE三个指标怎么配合看RankIQA 论文和大多数 IQA 工作都以 SROCC 作为主指标但工程落地时只看一个指标会误判指标全称关心什么什么时候看它SROCCSpearman 秩相关系数预测与 MOS 的单调一致性无参考场景、混合失真类型时LCCPearson 线性相关系数预测与 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 项目本身就是一条很扎实的基线跑通一次你对排序学习、迁移训练和图像质量评估的完整认知就建立起来了。希望帮到你。本文还有配套的精品资源点击获取