MLBench/ReaLens

1import argparse2from pathlib import Path3from util import util4import torch5import models6import data7 8 9class BaseOptions:10    """This class defines options used during both training and test time.11 12    It also implements several helper functions such as parsing, printing, and saving the options.13    It also gathers additional options defined in <modify_commandline_options> functions in both dataset class and model class.14    """15 16    def __init__(self):17        """Reset the class; indicates the class hasn't been initailized"""18        self.initialized = False19 20    def initialize(self, parser):21        """Define the common options that are used in both training and test."""22        # basic parameters23        parser.add_argument("--dataroot", required=True, help="path to images (should have subfolders trainA, trainB, valA, valB, etc)")24        parser.add_argument("--name", type=str, default="experiment_name", help="name of the experiment. It decides where to store samples and models")25        parser.add_argument("--checkpoints_dir", type=str, default="./checkpoints", help="models are saved here")26        # model parameters27        parser.add_argument("--model", type=str, default="cycle_gan", help="chooses which model to use. [cycle_gan | pix2pix | test | colorization]")28        parser.add_argument("--input_nc", type=int, default=3, help="# of input image channels: 3 for RGB and 1 for grayscale")29        parser.add_argument("--output_nc", type=int, default=3, help="# of output image channels: 3 for RGB and 1 for grayscale")30        parser.add_argument("--ngf", type=int, default=64, help="# of gen filters in the last conv layer")31        parser.add_argument("--ndf", type=int, default=64, help="# of discrim filters in the first conv layer")32        parser.add_argument("--netD", type=str, default="basic", help="specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator")33        parser.add_argument("--netG", type=str, default="resnet_9blocks", help="specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]")34        parser.add_argument("--n_layers_D", type=int, default=3, help="only used if netD==n_layers")35        parser.add_argument("--norm", type=str, default="instance", help="instance normalization or batch normalization [instance | batch | none | syncbatch]")36        parser.add_argument("--init_type", type=str, default="normal", help="network initialization [normal | xavier | kaiming | orthogonal]")37        parser.add_argument("--init_gain", type=float, default=0.02, help="scaling factor for normal, xavier and orthogonal.")38        parser.add_argument("--no_dropout", action="store_true", help="no dropout for the generator")39        # dataset parameters40        parser.add_argument("--dataset_mode", type=str, default="unaligned", help="chooses how datasets are loaded. [unaligned | aligned | single | colorization]")41        parser.add_argument("--direction", type=str, default="AtoB", help="AtoB or BtoA")42        parser.add_argument("--serial_batches", action="store_true", help="if true, takes images in order to make batches, otherwise takes them randomly")43        parser.add_argument("--num_threads", default=4, type=int, help="# threads for loading data")44        parser.add_argument("--batch_size", type=int, default=1, help="input batch size")45        parser.add_argument("--load_size", type=int, default=286, help="scale images to this size")46        parser.add_argument("--crop_size", type=int, default=256, help="then crop to this size")47        parser.add_argument("--max_dataset_size", type=int, default=float("inf"), help="Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.")48        parser.add_argument("--preprocess", type=str, default="resize_and_crop", help="scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]")49        parser.add_argument("--no_flip", action="store_true", help="if specified, do not flip the images for data augmentation")50        parser.add_argument("--display_winsize", type=int, default=256, help="display window size for both visdom and HTML")51        # additional parameters52        parser.add_argument("--epoch", type=str, default="latest", help="which epoch to load? set to latest to use latest cached model")53        parser.add_argument("--load_iter", type=int, default="0", help="which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]")54        parser.add_argument("--verbose", action="store_true", help="if specified, print more debugging information")55        parser.add_argument("--suffix", default="", type=str, help="customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}")56        # wandb parameters57        parser.add_argument("--use_wandb", action="store_true", help="if specified, then init wandb logging")58        parser.add_argument("--wandb_project_name", type=str, default="CycleGAN-and-pix2pix", help="specify wandb project name")59        self.initialized = True60        return parser61 62    def gather_options(self):63        """Initialize our parser with basic options(only once).64        Add additional model-specific and dataset-specific options.65        These options are defined in the <modify_commandline_options> function66        in model and dataset classes.67        """68        if not self.initialized:  # check if it has been initialized69            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)70            parser = self.initialize(parser)71 72        # get the basic options73        opt, _ = parser.parse_known_args()74 75        # modify model-related parser options76        model_name = opt.model77        model_option_setter = models.get_option_setter(model_name)78        parser = model_option_setter(parser, self.isTrain)79        opt, _ = parser.parse_known_args()  # parse again with new defaults80 81        # modify dataset-related parser options82        dataset_name = opt.dataset_mode83        dataset_option_setter = data.get_option_setter(dataset_name)84        parser = dataset_option_setter(parser, self.isTrain)85 86        # save and return the parser87        self.parser = parser88        return parser.parse_args()89 90    def print_options(self, opt):91        """Print and save options92 93        It will print both current options and default values(if different).94        It will save options into a text file / [checkpoints_dir] / opt.txt95        """96        message = ""97        message += "----------------- Options ---------------\n"98        for k, v in sorted(vars(opt).items()):99            comment = ""100            default = self.parser.get_default(k)101            if v != default:102                comment = "\t[default: %s]" % str(default)103            message += "{:>25}: {:<30}{}\n".format(str(k), str(v), comment)104        message += "----------------- End -------------------"105        print(message)106 107        # save to the disk108        expr_dir = Path(opt.checkpoints_dir) / opt.name109        util.mkdirs(expr_dir)110        file_name = expr_dir / f"{opt.phase}_opt.txt"111        with open(file_name, "wt") as opt_file:112            opt_file.write(message)113            opt_file.write("\n")114 115    def parse(self):116        """Parse our options, create checkpoints directory suffix, and set up gpu device."""117        opt = self.gather_options()118        opt.isTrain = self.isTrain  # train or test119 120        # process opt.suffix121        if opt.suffix:122            suffix = ("_" + opt.suffix.format(**vars(opt))) if opt.suffix != "" else ""123            opt.name = opt.name + suffix124 125        self.print_options(opt)126        self.opt = opt127        return self.opt128