本文介绍RealBasicVSR论文的复现,不对论文进行解读,大家可以去参考其他大神的博客,这里只讲如何去实现。本来想基于MMEditing去写,突然发现MMEditing库已经更新了,改名为MMagic了,所以这里重新整理一下。

本文主要介绍如何使用MMagic在自己的视频超分数据集上进行RealBasicVSR训练和推理,平台为Ubuntu20.04,显卡为NVIDIA RTX 3090 24G X4。

引流:BasicVSR,BasicVSR++,RealBasicVSR,EDVR,MMagic

1. 下载代码仓库

git clone https://github.com/open-mmlab/mmagic.git

2. 安装

首先,根据官方引导安装好PyTorch。

其次,安装MIM。

pip安装,或者pip3安装都可以,不过还是建议从源码安装mmagic,本文使用源码安装的方式。

pip3 install openmimmim install 'mmcv>=2.0.0'mim install 'mmengine'mim install 'mmagic'# or pip install openmimmim install 'mmcv>=2.0.0'mim install 'mmengine'mim install 'mmagic'
cd mmagicpip install -e .

安装好之后,就可以开始准备训练了。

3. 数据集准备

观察mmagic文件目录,

和之前的mmediting其实非常相似,首先观察mmagic/datasets目录,可以看到很多基准数据集加载的实现。

我们直接从腾讯的一个比赛拿到了数据集,如果大家对REDS数据集很熟悉的话,很多视频超分的数据集文件组织形式就可以按照这个来做,不用自己重新实现一个加载类,我们这次用到的就是这个mmagic/datasets/basic_image_dataset.py类,里面已经实现了如何加载。

里面实现了很多方法,但是最重要的还是load_data_list(),这个和PyTorch的Dataset中的__getitem__()是一样的,如果你的数据集目录也是按照REDS组织的,这些基本不用修改。

我们的数据集组织形式如下:

分为低分辨率目录(LR)和对应的高分辨率目录(GT),每个目录下的片段命名为”{:03d}”的形式,每个片段下的视频帧命名为”{:08d}.png”的形式。

如果你觉得用人家这个已经写好的Dataset不方便,那么完全可以按照类似的形式重新写一个,举个例子,我想针对自己的数据集重新写一个类,我取名为tencent_sdr_dataset.py,放在和basic_image_dataset.py同级目录下:

里面具体的实现和没什么区别,我只改了Class的命名,命名为TencentSDRDataset:

然后要用到这个Class,需要在datasets目录下的__init__.py目录里修改添加:

这些都修改好后就可以使用这个Class了。

修改好Datasets部分后,可以考虑修改配置文件了,打开RealBasicVSR的配置目录configs/real_basicvsr目录,并新建文件.py脚本文件,里面的详细内容可以从其他配置里先拷贝,然后再做修改。

我使用的配置详细如下:

图中红框是在原配置的基础上进行修改的部分。首先修改一下自己的实验目录命名。其次,修改数据集路径为自己的路径,同时我自己添加了validation_root当做验证数据集路径,我才用的是REDS数据集的验证集。然后,在train_dataloader里,将dataset的type修改为我们之前的Class命名即可,metainfo里面的dataset_type修改一下是没有关系的,我改为了tencent作为区分,task_name保持不变,data_prefix是训练数据集路径的前缀,img和gt分别存放LQ和GT。

这里剩余的配置和前面一样,修改了val_dataloader里面的相应内容。因为我们的训练集有点大,所以修改了max_iters增加训练次数。

这些都修改好后,就可以开始训练了,我使用的是多卡训练,直接调用官方的命令即可。

./tools/dist_train.sh configs/real_basicvsr/realbasicvsr_wogan-c64b20_4K.py 4

这里就是调用自己的配置文件,后面的是GPU数量,我用4张卡训练。

训练的默认目录就在work_dirs下面,还可以看到你训练的log文件以及保存的权重。这样,用RealBasicVSR在你自己的数据集上训练也就ok啦。

模型的推理并不复杂,大家参考官方的inference文档即可。至于如何添加模型,添加损失,大家参考我之前写的这篇BasicVSR复现即可,mmagic下的修改方式是差不多的。