train_restoration.py

import argparse
import os
import yaml
import torch
import torch.utils.data
import numpy as np
import datasets
from models import DenoisingDiffusion_Restoration


def parse_args_and_config():
    parser = argparse.ArgumentParser(description='Training Patch-Based Denoising Diffusion Models')
    parser.add_argument("--config", type=str, default='Restoration.yml', help="Path to the config file")
    parser.add_argument("--phase", type=str, default='test', help="val(generation)")    
    parser.add_argument("--data_type", type=str, default='VI', help="data type options: ['LOL', 'IR', 'VI', 'CT', 'MRI', 'CT_norm', MRI_norm']")
    parser.add_argument('--name', type=str, default='Restoration_VI_MSRS', help='folder name to save outputs')
    parser.add_argument("--data_dir", type=str, default='data/IVIF_degraded/train', help="root data path")  
    parser.add_argument('--resume', default='', type=str, help='Path for checkpoint to load and resume')
    parser.add_argument("--sampling_timesteps", type=int, default=20, help="Number of implicit sampling steps for validation image patches")
    parser.add_argument('--seed', default=61, type=int, metavar='N', help='Seed for initializing training (default: 61)')     
    parser.add_argument('-gpu', '--gpu_ids', type=str, default="0")
    args = parser.parse_args()

    with open(os.path.join("configs", args.config), "r") as f:
        config = yaml.safe_load(f)
    new_config = dict2namespace(config)
    return args, new_config

def dict2namespace(config):
    namespace = argparse.Namespace()
    for key, value in config.items():
        if isinstance(value, dict):
            new_value = dict2namespace(value)
        else:
            new_value = value
        setattr(namespace, key, new_value)
    return namespace


def main():
    args, config = parse_args_and_config()
    # setup device to run
    device = torch.device("cuda:{}".format(args.gpu_ids)) if torch.cuda.is_available() else torch.device("cpu")
    print("Using device: {}".format(device))
    config.device = device

    # set random seed
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
    torch.backends.cudnn.benchmark = True

    # data loading
    print("=> using dataset '{}'".format(config.data.dataset))
    DATASET = datasets.__dict__[config.data.dataset](config, args)

    # create model
    print("=> creating denoising-diffusion model...")
    diffusion = DenoisingDiffusion_Restoration(args, config)
    diffusion.train(DATASET)


if __name__ == "__main__":
    main()