The overall code framework is shown in the following figure. It mainly consists of fice parts - Config, Data, DataOps, Model, and Network.
For example, once train.py is called (i.e. python train.py -opt options/sr/train_sr.yml
), a sequence of actions will follow this command:
- Config - Reads the configuration from /options/sr/train_sr.yml which is a YAML file. Then passes the configuration values from it down through the following steps.
- Data - Creates the train and validation dataloaders.
- DataOps - Contains most of the data operations, specially related to the different augmentation options and functions used in the dataloaders, losses and other components.
- Model - Creates the chosen model.
- Network - Creates the chosen network.
- Finally train.py - Starts to train the model. Other actions like logging, saving intermediate models, validation, updating learning rate, e.t.c. are also done during training.
Moreover, there are also Utilities and Useful script available to use for various operations, like configuring your dataset.
/options Configure the options for data loader, network structure, losses, training strategies and other components
- YAML (or JSON) files are used to configure options and /options/options.py will convert the file to a python dictionary (where missing keys return
None
instead ofException
). - YAML files use
~
forNone
; and officially supports comments with#
. (JSON files usenull
forNone
; and supports//
comments instead). - Supports
debug
mode, i.e, if the model name containsdebug
, it will trigger debug mode. - The configuration file and descriptions can be found in /options.
/data A data loader to provide data for training, validation, and testing
- A separate data loader module. You can modify/create data loader to meet your own needs. The data loaders are constructed in /data/__init__.py.
- By default, the package uses cv2 package to do image processing, which provides rich operations. PIL can also be used with a subset of options. However, it is possible to integrate other options if required.
- Supports reading files from image folder or lmdb file. For faster IO during training (even if on an SSD) it is recommended to create and use an lmdb dataset if possible.used.
- The standard tensor format used is NCHW, [0,1], RGB, torch float tensor, but they can be normalized using the
znorm
option if needed.
/dataops Data operations options
/dataops hosts a large variety of data operations, like: filters, colors, images augmentation, batch augmentations, differential augmentations, and others that can be used at different points of the architecture when needed.
- /dataops/common.py provides useful tools and functions that are used in other modules, such as image reading and saving, conversion between numpy arrays and PyTorch tensors and others.
- /dataops/imresize.py contains a resizing algorithm that supports antialiasing and solves the incorrectness issues from resizing with standard frameworks (OpenCV, PIL, PyTorch, TensorFlow, others) and a vast array of interpolation methods.
- The directory also contains the augmennt subtree for all the image transforms used in the dataloader pipelines.
/models Construct different models for training and testing
- The chosen model is constructed in /models/__init__.py according to selected model type.
- A model mainly consists of two parts - a Network structure, and a Model definition (e.g. loss configuration,optimization step, e.t.c.).
- Based on the /models/base_model.py, we define different models, e.g., SR, SFTGAN_ACD, VSR.
- SR is used to train super resolution (ESRGAN, PAN, SRGAN, etc) or (Super) Restoration (denoising, deblurring, etc) models, but based on it, additional models that require specific parameters optimization strategy or options can be defined, such as PPON.
/models/modules Construct different network architectures
- The network is constructed in /models/network.py and the detailed structures are in /models/modules/architectures.
- Some useful blocks can be found in /models/modules/architectures/block.py, which in general are common to multiple networks and it is flexible to construct your network structures with these pre-defined blocks.
- You can also easily write your own network architecture in a separate file like /models/modules/architectures/sft_arch.py or /models/modules/architectures/PAN_arch.py.
/utils Provides useful utilities
- /utils/util.py provides logging service during training and testing.
- /utils/progress_bar.py provides a progress bar which can print the progress.
- Support to use tensorboard to visualize and compare training loss, validation PSNR, SSIM, LPIPS, e.t.c.