使用自定义数据集微调Real-ESRGAN图片超分模型

1 为什么要微调Real-ESRGAN模型

Real-ESRGAN github仓库放出的开源模型基本上都是基于开源数据集进行训练的，比如DIV2K、Flickr2K等等，在一些特殊场景下其表现得并不是很好，比如商品图超分、书画超分等等，所以我们需要在自己的数据集上微调Real-ESRGAN的预训练模型，使其能适应某个特殊场景。

我们这里微调不采用LR-HR配对数据集，而是直接按照Real-ESRGAN的做法准备HR数据集，然后使用Real-ESRGAN的图片退化方法去做。

2 自定义数据集

在本文中，以LIU4K-v2这个4K 超分数据集为例，其包含了1600张高分辨率训练集图片和400张验证集图片，其图像分辨率至少为

3K，大部分为4K-6K，且包含了非常多样化的背景和物体，并且低级信号分布也多样而复杂，涵盖动物、建筑、风景、街景等等。

数据集官方链接：https://structpku.github.io/LIU4K_Dataset/LIU4K_v2.html。

我们的微调数据集以LIU4K-v2的train+val的2000张图片为微调数据集，将所有图片重命名为%07d.png的形式，比如0000001.png ~ 0002000.png。

3 数据集处理

按照Real-ESRGAN github仓库的微调教程 主要有以下几个步骤。

3.1 生成多尺寸图片（可选步骤）

对第2节中的HR图像进行下采样，获取多尺度的标准参考（Ground-Truth）图像。

python scripts/generate_multiscale_DF2K.py --input 数据集原始高清图片输入文件夹路径 --output 数据集多尺寸图片输出文件夹路径

示例命令：

python scripts/generate_multiscale_DF2K.py --input datasets/LIU4K-v2/LIU4K-v2_HR --output datasets/LIU4K-v2/LIU4K_multiscale

3.2 裁切为子图像（可选步骤）

这里我们不进行处理，主要是磁盘空间足够。

3.3 生成数据集元信息文件

这个元信息文件主要是包含了数据集图片的路径信息的txt文件。

（1）如果你跳过了3.1节中生成多尺寸图片，那么meta_info可通过以下命令生成

 python scripts/generate_meta_info.py --input 数据集原始高清图片输入文件夹路径 --root 数据集根目录路径 --meta_info 数据集元信息文件保存路径

示例命令：

 python scripts/generate_meta_info.py --input datasets/LIU4K-v2/LIU4K-v2_HR --root datasets/LIU4K-v2 --meta_info datasets/LIU4K-v2/meta_info/meta_info_LIU4K-v2.txt

（2）如果你在3.1节中生成了多尺寸图片，那么meta_info可通过以下命令生成

 python scripts/generate_meta_info.py --input 数据集原始高清图片输入文件夹路径, 数据集多尺寸图片输出文件夹路径 --root 数据集根目录路径, 数据集根目录路径 --meta_info 数据集元信息文件保存路径

示例命令：

 python scripts/generate_meta_info.py --input datasets/LIU4K-v2/LIU4K-v2_HR, datasets/LIU4K-v2/LIU4K_multiscale --root datasets/LIU4K-v2, datasets/LIU4K-v2 --meta_info datasets/LIU4K-v2/meta_info/meta_info_LIU4K-v2_multiscale.txt

4 Real-ESRGAN模型微调

4.1 下载预训练模型

从Real-ESRGAN的github仓库下载RealESRGAN_x4plus.pth和RealESRGAN_x4plus_netD.pth两个预训练模型文件放到experiments/pretrained_models目录下。

4.2 修改配置文件

拷贝一份options/finetune_realesrgan_x4plus.yml文件，重名为options/finetune_realesrgan_x4plus_liu4k_v2.yml

修改主要内容如下，
（1）修改实验名称

# general settings
name: finetune_RealESRGANx4plus_400k_LIU4kv2 # 修改实验名称

(2) 修改微调数据集路径

# general settings
name: finetune_RealESRGANx4plus_400k_LIU4kv2 # 修改实验名称
model_type: RealESRGANModel
scale: 4
num_gpu: auto
manual_seed: 0

......

datasets:
  train:
    name: LIU4k-v2 # 修改名称
    type: RealESRGANDataset
    dataroot_gt: ./datasets/LIU4K-v2 # 修改为你的数据集文件夹根目录
    meta_info: ./datasets/LIU4K-v2/meta_info/meta_info_LIU4K-v2_multiscale.txt # 修改为你自己生成的元信息txt
    io_backend:
      type: disk

（3）修改预训练模型
默认的微调配置finetune_realesrgan_x4plus.yml中的pretrain_network_g的权重指定的是RealESRNet_x4plus.pth

path:
  # use the pre-trained Real-ESRNet model
  pretrain_network_g: experiments/pretrained_models/RealESRNet_x4plus.pth

我们需要直接在RealESRGAN_x4plus.pth上进行微调，所以这里修改为

path:
  # use the pre-trained Real-ESRNet model
  pretrain_network_g: experiments/pretrained_models/RealESRGAN_x4plus.pth

（4）修改验证数据集路径和验证参数

  # Uncomment these for validation # 这里可以增加你的验证数据集
  # val:
  #   name: validation
  #   type: PairedImageDataset
  #   dataroot_gt: path_to_gt
  #   dataroot_lq: path_to_lq
  #   io_backend:
  #     type: disk

  # Uncomment these for validation # 这里可以增加你的数据集验证设置
 # validation settings
 # val:
 #   val_freq: !!float 5e3
 #   save_img: True

 #   metrics:
 #     psnr: # metric name
 #       type: calculate_psnr
 #       crop_border: 4
 #       test_y_channel: false

配置文件的其他超参数可以根据需求自行修改，比如预训练模型文件保存频率save_checkpoint_freq等参数。

4.3 模型微调

单个gpu模型微调可使用以下命令，--auto_resume用于断点续训

python realesrgan/train.py -opt options/finetune_realesrgan_x4plus_liu4k_v2.yml --auto_resume

多个gpu模型微调可使用以下命令

CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/finetune_realesrgan_x4plus_liu4k_v2.yml.yml --launcher pytorch --auto_resume

微调的模型和相关信息保存在experiments/pretrained_models目录的finetune_xxx文件夹下。

4.4 使用自己微调的模型进行推理

使用以下命令基于自己的微调模型进行推理

python inference_realesrgan.py -n RealESRGAN_x4plus -i inputs --outscale 4 --model_path 你自己微调模型的路径

参考链接

• https://github.com/xinntao/Real-ESRGAN
• https://zhuanlan.zhihu.com/p/674911337
• https://rockyshikoku.medium.com/train-realesrgan-with-your-own-data-to-create-customized-models-72672dbb7492