从零构建BasicVSR实战基于MMEditing的视频超分辨率全流程解析与REDS4深度评测最近在折腾视频超分辨率项目发现不少朋友对BasicVSR这个CVPR 2021的经典工作很感兴趣但实际动手时总在环境配置、数据准备和MMEditing框架使用上遇到各种“坑”。今天我就把自己从环境搭建到模型测试、再到自定义训练的全过程梳理出来希望能帮你少走弯路。BasicVSR作为循环结构的视频超分网络在保持轻量化的同时实现了出色的时间一致性恢复特别适合处理动态场景。不过官方代码虽然开源但想要在自己的机器上顺利跑起来还是需要一些工程上的技巧。这篇文章不会重复论文里的公式推导而是聚焦于如何把理论变成可运行的代码我会分享具体的配置文件修改、数据组织规范以及调试过程中那些文档里没写的细节。1. 环境准备与MMEditing框架深度解析在开始复现任何模型之前搭建一个稳定、兼容的开发环境是第一步。MMEditing作为OpenMMLab旗下的图像与视频编辑工具箱封装了大量经典和前沿的模型其模块化设计既带来了灵活性也增加了初学者的学习成本。1.1 依赖安装与版本锁定我强烈建议使用Conda来管理Python环境避免与系统Python产生冲突。以下是我验证过的一组兼容版本特别是PyTorch与CUDA的匹配至关重要。# 创建并激活专用环境 conda create -n mmediting python3.8 -y conda activate mmediting # 安装PyTorch请根据你的CUDA版本选择这里以CUDA 11.3为例 conda install pytorch1.11.0 torchvision0.12.0 torchaudio0.11.0 cudatoolkit11.3 -c pytorch # 安装MMCVMMEditing的基础 pip install mmcv-full1.6.0 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.11.0/index.html # 克隆MMEditing仓库并安装 git clone https://github.com/open-mmlab/mmediting.git cd mmediting pip install -v -e . # “-e”代表可编辑模式方便修改源码注意MMCV-full的版本必须与PyTorch和CUDA版本严格对应。如果安装失败可以去OpenMMLab官网查看详细的版本兼容性表格。安装完成后可以通过一个简单的导入测试来验证import mmedit print(mmedit.__version__) import torch print(torch.__version__, torch.cuda.is_available())1.2 理解MMEditing的项目结构很多新手拿到代码库后不知从何下手主要是因为不熟悉其目录结构。MMEditing的核心逻辑都藏在configs和mmedit这两个目录里。mmediting/ ├── configs/ # 所有模型的配置文件 │ └── restorers/ # 图像/视频修复模型BasicVSR在此 │ └── basicvsr/ │ ├── basicvsr_reds4.py │ └── basicvsr_vimeo90k.py ├── mmedit/ # 核心源代码 │ ├── models/ # 模型定义、损失函数 │ ├── datasets/ # 数据加载与处理管道 │ ├── core/ # 训练、测试、评估的流程控制 │ └── apis/ # 高级API接口 └── tools/ # 训练、测试、分析的脚本入口 ├── train.py ├── test.py └── dist_train.py # 分布式训练脚本关键认知在MMEditing中几乎所有的实验设置模型结构、数据集、训练策略、优化器都通过一个Python格式的配置文件.py文件来定义。你不需要也不建议直接修改源代码而是通过继承和修改这些配置文件来定制你的实验。这就像搭积木官方提供了各种标准件_base_配置你只需要替换自己需要的部分。2. REDS4数据集准备与配置文件精讲REDS数据集是视频超分领域的常用基准而REDS4是其一个包含4个视频序列的轻量测试子集非常适合快速验证模型效果。2.1 数据集的正确组织方式从官网下载REDS数据集后你不能直接把一堆视频文件扔进去。MMEditing有严格的数据组织要求目的是让Dataset类能够自动索引到每一帧。你需要将数据集整理成如下结构data/ └── REDS/ ├── train_sharp/ # 训练集高清帧 │ ├── 000/ │ │ ├── 00000000.png │ │ ├── 00000001.png │ │ └── ... │ ├── 001/ │ └── ... ├── train_sharp_bicubic/ # 训练集下采样后的低清帧X4 │ ├── 000/ │ └── ... ├── val_sharp/ # 验证集高清帧REDS4在这里 │ ├── 000/ │ ├── 011/ │ ├── 015/ │ └── 020/ └── val_sharp_bicubic/ # 验证集低清帧 ├── 000/ ├── 011/ ├── 015/ └── 020/对于REDS4测试我们只关心val_sharp和val_sharp_bicubic下的那四个文件夹000, 011, 015, 020。每个文件夹里的PNG序列就是一段视频。2.2 解剖BasicVSR的配置文件现在来到最关键的一步修改配置文件。我们以configs/restorers/basicvsr/basicvsr_reds4.py为例。不要被它上百行的代码吓到我们只需关注几个关键部分。首先找到数据管道data配置部分。你需要修改所有指向数据集的路径。# 原始配置片段可能需要修改的地方 data dict( test_dataloaderdict( samples_per_gpu1, workers_per_gpu1, drop_lastFalse, shuffleFalse), testdict( typeSRFolderMultipleGTDataset, lq_folderdata/REDS/val_sharp_bicubic/X4, # 修改为你的低清帧路径 gt_folderdata/REDS/val_sharp, # 修改为你的高清帧路径 scale4, filename_tmpl{}, pipelinetest_pipeline))你需要将lq_folder和gt_folder的路径改为你本地存放REDS4数据的绝对路径或相对路径。例如如果你的数据放在/home/user/projects/data/REDS下那么路径就应该是/home/user/projects/data/REDS/val_sharp_bicubic/X4。其次加载预训练模型。BasicVSR的官方预训练权重可以在模型库找到。下载basicvsr_reds4_20120409-0e599677.pth文件。在配置文件中通常会有load_from参数或者在测试时通过命令行指定。我更喜欢在配置文件中直接指定避免每次输入命令都带参数。# 在配置文件末尾或合适位置添加 load_from checkpoints/basicvsr_reds4_20120409-0e599677.pth # 你的权重文件路径提示SPyNet光流网络是BasicVSR的一个组件其预训练权重会在代码运行时自动下载如果本地没有。如果网络环境不好导致下载失败你可以手动从MMFlow的模型库下载spynet_20210409-c6c1bd09.pth并放在~/.cache/torch/hub/checkpoints/目录下。3. 模型测试与结果可视化全流程配置修改妥当后就可以开始测试了。测试的目的不仅是得到PSNR/SSIM数字更要直观地看到超分后的视频效果。3.1 单卡测试命令与参数解析打开终端进入MMEditing根目录运行以下命令python tools/test.py \ configs/restorers/basicvsr/basicvsr_reds4.py \ checkpoints/basicvsr_reds4_20120409-0e599677.pth \ --save-path results/basicvsr_reds4_demo这条命令分解开来tools/test.py: 测试脚本。第一个参数配置文件的路径。第二个参数模型权重文件的路径如果配置文件中已指定load_from这里可以省略。--save-path:极其重要的参数指定生成的高清帧保存目录。如果不设置默认只计算指标不保存图像。运行后你会看到终端输出每个视频序列的PSNR和SSIM以及平均值。同时在--save-path指定的目录下会生成和输入结构一致的文件夹里面就是超分辨率后的每一帧图像。3.2 结果分析与常见问题排查第一次运行很可能不会一帆风顺。下面是一些我踩过的坑和解决方案问题1KeyError: lq或KeyError: gt原因数据管道pipeline的某个环节期望在数据字典里找到lq低清或gt高清键但没找到。排查首先检查lq_folder和gt_folder路径是否正确并且文件夹内是否有PNG文件。其次检查filename_tmpl参数是否与你的文件名匹配。REDS数据集帧文件名是8位数字如00000000.png所以filename_tmpl{}是正确的。如果你的文件名是frame_001.png则需要改为filename_tmplframe_{}。问题2CUDA out of memory原因视频超分很耗显存BasicVSR默认会缓存多个相邻帧。解决在配置文件中调整测试参数。找到model配置部分尝试减小spynet_pretrained或调整padding。更直接的方法是修改test_dataloader中的samples_per_gpu确保其为1。问题3生成的图像全黑或全白原因图像张量的数值范围0-1或0-255与模型期望的不匹配或者后处理output管道有问题。排查检查配置文件中test_pipeline的最后几步。通常包含Collect收集键值、ImageToTensor转换、Normalize归一化和FormatShape。确保训练和测试的Normalize参数均值、标准差一致。一个快速验证的方法是用Matplotlib或OpenCV直接读取--save-path下的一张输出图像看看像素值是否在合理范围0-255。为了更直观地对比效果你可以写一个简单的脚本将低清LR、超分SR和真实高清HR帧拼接到一起import cv2 import numpy as np lr_img cv2.imread(data/REDS/val_sharp_bicubic/X4/000/00000000.png) sr_img cv2.imread(results/basicvsr_reds4_demo/000/00000000.png) hr_img cv2.imread(data/REDS/val_sharp/000/00000000.png) # 调整大小使其一致如果尺度不同 lr_img_resized cv2.resize(lr_img, (hr_img.shape[1], hr_img.shape[0]), interpolationcv2.INTER_CUBIC) concat_img np.concatenate([lr_img_resized, sr_img, hr_img], axis1) cv2.imwrite(comparison_000.png, concat_img)4. 进阶在自定义数据集上训练BasicVSR测试通过后如果你想在自己的数据上训练模型以适配特定场景如动画、监控视频就需要准备训练集并修改训练配置。4.1 准备训练数据与索引文件训练需要的数据结构类似REDS。此外MMEditing通常需要一个索引文件ann_file这是一个文本文件如.txt列出了所有用于训练的视频片段对。每一行格式通常为视频文件夹名称 起始帧索引 结束帧索引。例如你的训练数据有video1,video2两个文件夹每个有100帧你想以每30帧为一个片段那么索引文件可能如下video1 0 29 video1 30 59 video1 60 89 video2 0 29 video2 30 59 video2 60 89MMEditing在tools/data/目录下提供了制作各种数据集索引的脚本你可以参考super-resolution下的脚本来生成自己的。4.2 配置训练参数复制一份basicvsr_reds4.py并重命名如basicvsr_custom.py然后进行大幅修改修改数据配置将data字典下的train部分指向你的训练数据路径和索引文件。修改模型输入通道如果你的视频是RGB彩色保持默认3通道即可。如果是灰度或其它格式需要修改model中generator的mid_channels和spynet的输入通道。调整训练策略total_iters总迭代次数、lr_config学习率策略、optimizer优化器设置都需要根据你的数据集大小和硬件条件调整。对于小数据集可以减小total_iters增加验证频率val_interval。修改验证集将data中的val部分指向你的验证集可以先用REDS4做验证确保训练过程正常。一个简化版的训练配置修改示例如下# 在自定义配置文件中 _base_ ./basicvsr_reds4.py # 继承基础配置 # 覆盖数据配置 data dict( traindict( lq_folderdata/custom/train_lq, gt_folderdata/custom/train_hr, ann_filedata/custom/train_ann.txt, # 你的训练索引文件 num_input_frames10, # 输入帧数根据你的需求调整 pipelinetrain_pipeline), valdict( lq_folderdata/REDS/val_sharp_bicubic/X4, # 暂时用REDS4验证 gt_folderdata/REDS/val_sharp, pipelineval_pipeline), testdict( lq_folderdata/REDS/val_sharp_bicubic/X4, gt_folderdata/REDS/val_sharp, pipelinetest_pipeline)) # 调整训练策略 total_iters 300000 # 根据你的数据量调整 lr_config dict( policyStep, by_epochFalse, step[200000, 250000], gamma0.5) optimizer dict(typeAdam, lr2e-4, betas(0.9, 0.999)) # 工作目录用于保存日志和模型权重 work_dir ./work_dirs/basicvsr_custom_exp14.3 启动训练与监控使用分布式训练可以加速过程即使只有一块卡也推荐用分布式接口方便未来扩展CUDA_VISIBLE_DEVICES0,1 ./tools/dist_train.sh configs/restorers/basicvsr/basicvsr_custom.py 2这里CUDA_VISIBLE_DEVICES0,1指定使用第0和1号GPU最后的2表示使用2个GPU进程。训练开始后在work_dir下会生成日志文件和模型权重快照。你可以使用TensorBoard来监控训练损失和验证指标tensorboard --logdir ./work_dirs/basicvsr_custom_exp1打开浏览器访问localhost:6006就能看到实时的训练曲线。4.4 自定义损失函数以Charbonnier Loss为例有时你可能想尝试论文中提到的其他损失函数。MMEditing的损失函数注册机制非常清晰。假设你想在L1损失基础上增加一个Charbonnier Loss一种鲁棒的L1损失变体。首先你不需要自己实现因为MMEditing已经内置了。你只需要在配置文件的model部分修改pixel_loss的类型即可。但让我们看看如何找到并确认它在mmedit/models/losses/pixelwise_loss.py中你可以找到L1Loss,MSELoss和CharbonnierLoss的定义。要使用它你需要在配置文件的loss_config中进行设置。不过BasicVSR的配置文件可能将损失函数定义在了model的train_cfg或直接作为pixel_loss参数。查找基础配置文件你会发现类似这样的结构model dict( typeBasicVSR, generatordict(...), pixel_lossdict(typeCharbonnierLoss, loss_weight1.0, reductionmean), ... )如果你想同时使用多个损失例如L1损失和感知损失VGG Loss可以定义一个ComposedLossmodel dict( typeBasicVSR, generatordict(...), pixel_lossdict( typeComposedLoss, loss_types[L1Loss, PerceptualLoss], loss_weights[1.0, 0.1]), ... )然后确保PerceptualLoss在mmedit/models/losses/中有定义并被正确注册。这种模块化的设计让你可以像搭积木一样组合不同的损失函数、模型组件和数据增强策略。整个流程走下来你会发现MMEditing虽然初期学习曲线陡峭但其严谨的设计模式一旦掌握就能极大地提升实验效率。从修改配置到训练新模型大部分工作都不需要触碰核心代码这保证了代码的整洁和可复现性。当然调试时直接阅读源码理解数据流和模型前向传播过程是解决问题的最终捷径。