options.py

import argparse


def get_message(parser, args):
    r"""
    References:
        1. https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/f13aab8148bd5f15b9eb47b690496df8dadbab0c/options/base_options.py#L88
    """
    message = ''
    message += '----------------- Options ---------------\n'
    for k, v in sorted(vars(args).items()):
        comment = ''
        default = parser.get_default(k)
        if v != default:
            comment = '\t[default: %s]' % str(default)
        message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
    message += '----------------- End -------------------'
    return message


def get_args():
    parser = argparse.ArgumentParser(description='RMEP')
    
    # basic parameters
    parser.add_argument('--dataroot', type=str, default='datasets', help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')
    parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
    parser.add_argument('--checkpoints_dir', type=str, default='checkpoints', help='models are saved here')
    parser.add_argument('--isTest', action='store_true', help='default is False (training), if specified, crop_size and preprocess will be set to be None')

    # model parameters
    parser.add_argument('--model_name', type=str, default='rmep', help='which model to use')
    parser.add_argument('--model_type', type=str, default='reg', help='type of the model: [reg | clf | clf_multi]')
    parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels: 3 for RGB and 1 for grayscale')
    parser.add_argument('--output_nc', type=int, help='# of output numbers, for classification and regression')
    parser.add_argument('--nrf', type=int, default=64, help='# of regression filters in the last conv layer')
    parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization [instance | batch | none]')
    parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal | xavier | kaiming | orthogonal]')
    parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
    parser.add_argument('--continue_train', action='store_true', help='train again if specified')
    parser.add_argument('--num_blk', type=int, default=6, help='number of resnet blocks')
    parser.add_argument('--pooling_type', type=str, default='max', help='type of pooling layer [max | avg]')

    # dataset parameters
    parser.add_argument('--num_workers', default=8, type=int, help='# threads for loading data')
    parser.add_argument('--batch_size', type=int, default=16, help='input batch size')
    parser.add_argument('--crop_size', type=int, default=480, help='then crop to this size, no use if the preprocess option is none')
    parser.add_argument('--preprocess', type=str, default='none', help='scaling and cropping of images at load time [crop_and_flip | flip | crop | none]')
    parser.add_argument('--HSI', action='store_true', help='the input images are hyperspectral images if specified')

    # additional parameters
    parser.add_argument('--epoch', type=str, default='best', help='which epoch to save/load?')
    parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')
    parser.add_argument('--start_epoch', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')
    parser.add_argument('--save_start_epoch', type=int, default=150, help='start epoch to save')
    parser.add_argument('--save_freq', type=int, default=10, help='save model every save_freq')

    # training parameters
    parser.add_argument('--epochs', type=int, default=50, help='number of epochs with the initial learning rate')
    parser.add_argument('--epochs_decay', type=int, default=150, help='number of epochs to linearly decay learning rate to zero')
    parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
    parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
    parser.add_argument('--output_params', nargs='+', help='output parameters')
    parser.add_argument('--lr_policy', type=str, default='linear', help='learning rate policy. [linear | step | plateau | cosine]')
    parser.add_argument('--seed', type=int, default=42, help='determines the random state for random, numpy, torch')
    args = parser.parse_args()

    return args, get_message(parser, args)