独自のデータセットでReal-ESRGANをファインチューニングする。

Real-ESRGANのモデルに対し追加学習（ファインチューニング）することで、特定の画像に特化したモデルを作成することが出来るとのことですので試してみたいと思います。

以下のページの手順通りに進める。

Real-ESRGAN/docs/Training.md at master · xinntao/Real-ESRGAN

Real-ESRGAN aims at developing Practical Algorithms for General Image/Video Restoration. - xinntao/Real-ESRGAN

Real-ESRGANのインストール
データセットの準備
トレーニングデータのダウンロード
設定ファイルの変更
ソースコードの修正
ファインチューンの実行
用意した画像
学習したモデルで拡大処理を実行
Google Colabでファインチューニング

Real-ESRGANのインストール

環境:
Ubuntu22.04がインストールされたPC
CPU:AMD Ryzen 5 7600 6-Core Processor
GPU:GeForce RTX 3060 (VRAM 12GB)
WSL-Ubuntuでも良いですが、学習でメインPCが占有される可能性が高いので専用PCを用意する。

git clone https://github.com/xinntao/Real-ESRGAN.git
cd Real-ESRGAN
mkdir venv
python -m venv venv
source venv/bin/activate
pip install basicsr
pip install facexlib
pip install gfpgan
pip install -r requirements.txt
python setup.py develop
deactivate

長時間対策してscreenコマンドの操作法を覚えて途中でsshを切断と再接続が出来るようにしておく。

Ubuntuで使うコマンド備忘録

使用したコマンドの備忘録です。ファイルをサブディレクトリを含めてzipファイルに圧縮（元ファイル削除)zip ps1-1.zip -m -r ./ps1「ps1.zip」は圧縮先のファイル「./ps1」が圧縮元のディレクトリ-mオプションで...

データセットの準備

高解像度画像をdatasets/DF2K/DIV2K_train_HR_subディレクトリにセット
低解像度画像をdatasets/DF2K/DIV2K_train_LR_bicubic_X4_subディレクトリにセット

generate_meta_info_pairdata.pyで設定ファイル生成

python scripts/generate_meta_info_pairdata.py --input datasets/DF2K/DIV2K_train_HR_sub datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub --meta_info datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt

設定ファイルの出力先datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt

トレーニングデータのダウンロード

wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models

wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth -P experiments/pretrained_models

-Pがダウンロード先のディレクトリですのでexperiments/pretrained_modelsにダウンロードされる。

設定ファイルの変更

options/finetune_realesrgan_x4plus_pairdata.yml をディレクトリに合わせて変更

train:
    name: DIV2K
    type: RealESRGANPairedDataset
    dataroot_gt: datasets/DF2K  # modify to the root path of your folder
    dataroot_lq: datasets/DF2K  # modify to the root path of your folder
    meta_info: datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt  # modify to your own generate meta info txt
    io_backend:
        type: disk

dataroot_gtが高解像度データセットを保存するディレクトリのルートdatasets/DF2K
dataroot_lqが低解像度データセットを保存するディレクトリのルートdatasets/DF2K
となっていますが、環境に合わせて変更しなさいとコメントがあるので要確認
meta_infoがdatasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt先ほど生成したファイルのパスを指定するようです。

pretrain_network_g: experiments/pretrained_models/RealESRNet_x4plus.pth

を

pretrain_network_g: experiments/pretrained_models/RealESRGAN_x4plus.pth

へ変更

このままだと筆者の環境だとVRAM不足となるのでbatch_size_per_gpuのサイズを小さな値にしました。

Real-ESRGANの追加学習で「CUDA out of memory」が発生した話

縮小などで劣化してしまった画像を綺麗な形に拡大処理してくれるAIの一つにReal-ESRGANがあります。学習結果であるモデルを差し替えることで画像の内容にに合わせた処理を行うことが出来るのですが、このモデルを追加学習することで自分好みにカ...

ソースコードの修正

実行すると

ModuleNotFoundError: No module named 'torchvision.transforms.functional_tensor'

が発生
エラーメッセージが発生したdegradations.pyをテキストエディタで

from torchvision.transforms.functional_tensor import rgb_to_grayscale

から

from torchvision.transforms.functional import rgb_to_grayscale

に変更

ファインチューンの実行

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

-optにoptions/finetune_realesrgan_x4plus_pairdata.ymlのパスがセットされています。
–auto_resumeオプションは字ずらからすると自動的に前回の続きから実行してくれるように見える。

用意した画像

高解像度は256×256サイズ、低解像度は64×64サイズと4倍に拡大することを想定してみました。

どうやって、データセットを用意するのかが問題です。

学習したモデルで拡大処理を実行

python inference_realesrgan.py -n RealESRGAN_x4plus -i inputs --outscale 4 --model_path experiments/finetune_RealESRGANx4plus_400k_pairdata/models/net_g_latest.pth

トーン処理された画像とグレースケールの画像を学習させたモデルを使うと以下のような結果になりました。

トーン処理された元画像

学習したモデルで拡大処理⇒GIMPで元画像と同じサイズに縮小

超解像による拡大することが本来の機能ですが、トーンの除去に使えないかと思い試してみました。

Google Colabでファインチューニング

ローカル環境では処理能力不足で学習時間が長時間になりがちなのでColabでもファインチューニングする環境を構築してみました。

2025/03/19日現在、上記と同じ手順で行うとColabでもエラーメッセージが表示されて詰まります。
ただ、Colabの場合上記のソースコードの変更が出来ないのでインストールするパッケージのバージョンを古い物にすることで対応します。

!pip install basicsr
!pip install facexlib
!pip install gfpgan
!pip install torch==2.0.1 torchvision==0.15.2 --extra-index-url https://download.pytorch.org/whl/cu118
!pip install numpy==1.26.4
!pip install -r requirements.txt
!python setup.py develop

requirements.txtは以下のように書き換えておく

basicsr>=1.4.2
facexlib>=0.2.5
gfpgan>=1.3.5
#numpy
opencv-python
Pillow
#torch>=1.7
#torchvision
tqdm

Colabはセッション時間が決められており、長時間の学習には向かないと思いますが、Real-ESRGANはauto_resumeオプションで前回の続きから実行することが出来ます。
Colabではきちんとレジュームされるか確認できていませんが、ローカル環境では機能していました。（数日に分けて実行しました。）

最後に時間経過でpipなどでインストールするパッケージのバージョンは上がるので、この手順でも実行できない日が来るでしょう。
Real-ESRGANのソースコードを新しいパッケージバージョンに合わせて修正が出来ればよいのですが、そうでなければローカル環境で実行環境を仮想化してバージョンを固定するしかないのかもしれません。