Deep learning training framework for image super resolution and restoration.
APACHE-2.0 License
A modern community fork of BasicSR and traiNNer-redux.
git clone https://github.com/the-database/traiNNer-redux.git
install.bat
, and for Linux users, from the terminal in the traiNNer-redux folder run chmod +x install.sh && ./install.sh
to install all Python dependencies to a new virtual environment.Refer to the wiki for a full training guide and benchmarks.
The repository comes with several configs that are ready to use out of the box, as well as a tiny dataset for testing purposes only. To confirm that your PC can run the training software successfully, run the following command from the traiNNer-redux
folder:
venv\Scripts\activate
python train.py --auto_resume -opt ./options/train/SPAN/SPAN.yml
You should see the following output within a few minutes, depending on your GPU speed:
...
2024-07-02 21:40:56,593 INFO: Model [SRModel] is created.
2024-07-02 21:40:56,668 INFO: Start training from epoch: 0, iter: 0
2024-07-02 21:41:17,816 INFO: [4x_SP..][epoch: 0, iter: 100, lr:(1.000e-04,)] [performance: 4.729] [eta: 14:11:33] l_g_mssim: 1.0000e+00 l_g_percep: 3.5436e+00 l_g_hsluv: 4.3935e-01 l_g_gan: 2.4346e+00 l_g_total: 7.4175e+00 l_d_real: 2.4136e-01 out_d_real: 2.9309e+00 l_d_fake: 5.2773e-02 out_d_fake: -2.4104e+01
The last line shows the progress of training after 100 iterations. If you get this far without any errors, your PC is able to train successfully. Press ctrl+C
to end the training run.
traiNNer-redux/options/train
, select the architecture you want to train, and open the yml
file in that folder in a text editor. A text editor that supports YAML syntax highlighting is recommended, such as VS Code or Notepad++. For example, to train SPAN, open traiNNer-redux/options/train/SPAN/SPAN.yml
.name
to the name of the model you want to train. Give it a unique name so you can differentiate it from other training runs.# 2, 4
means the architecture only supports a scale of 2 or 4.dataroot_gt
and dataroot_lq
under the train:
section. The HR images and LR images should match in numer of images and filenames. For each matching LR/HR pair, the image resolutions should work with the selected scale, so if a scale of 2 is selected then each HR must be 2x the resolution of its matching LR image.val_enabled
to true
and set the paths to your validation HR and LR images, at dataroot_gt
and dataroot_lq
under the val
section.pretrain_network_g
and remove the #
to uncomment that line.Run the following command to start training. Change ./options/train/arch/config.yml
to point to the config file you set up in the previous step.
venv\Scripts\activate
python train.py --auto_resume -opt ./options/train/arch/config.yml
For example, to train with the SPAN config:
venv\Scripts\activate
python train.py --auto_resume -opt ./options/train/SPAN/SPAN.yml
To pause training, press ctrl+C
or close the command window. To resume training, run the same command that was used to start training. The --auto_resume
flag will resume training from when it was paused.
Models are saved in the safetensors
format to traiNNer-redux/experiments/<name>/models
, where name
is whatever was used in the config file. chaiNNer can be used to run most models. If you want to run the model on images during training to monitor the progress of the model, set up validation in the config file, and find the validation results in traiNNer-redux/experiments/<name>/visualization
.
The test script can also be used to test trained models, which is required to test models with architectures not yet supported by chaiNNer. For example, to test SPANPlus model, open the config file at ./options/test/SPANPlus/SPANPlus.yml
, and update the following:
dataroot_lq
option to point to a folder that contains the images you want to run the model on.network_g
match the options under network_g
in the training config file that you used. For example, if you trained SPANPlus_STS
, then set the type to SPANPlus_STS
under network_g
in the test config file as well.pretrain_network_g
to point to the path of the model you want to test.Then run this command to run the model on the images as specified in the config file:
venv\Scripts\activate
python test.py -opt ./options/test/SPANPlus/SPANPlus.yml
If you want to convert your PyTorch models to ONNX format, you can use the convert_to_onnx.py
script to do so. First install the additional dependencies for ONNX:
venv\Scripts\activate
pip install .[onnx]
Then open a config file corresponding to the architecture of the model you trained. For example, if you trained SPANPlus
, open the config file at ./options/onnx/SPANPlus/SPANPlus.yml
, and update the following:
name
to the name of your model, the ONNX filename will include this name.network_g
match the settings you used to train your model.pretrain_network_g
to point to the path of your .safetensors
or .pth
model that you want to convert.onnx
section of the config file as needed.Then run this command to do the conversion (make sure the path points to the .yml
file you edited):
venv\Scripts\activate
python convert_to_onnx.py -opt ./options/onnx/SPANPlus/SPANPlus.yml
The converted onnx files will be saved to the onnx
directory.
Please see the contributing page for more info on how to contribute.
traiNNer-redux is released under the Apache License 2.0. See LICENSE for individual licenses and acknowledgements.