|
- import argparse
- import os
- import torch
- import model
- import data
- from util import util
- import sys
-
-
- class BaseOptions():
- def __init__(self):
- self.parser = argparse.ArgumentParser()
- self.initialized = False
-
- def initialize(self, parser):
- # base define
- parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment.')
- parser.add_argument('--model', type=str, default='rec', help='name of the model type.')
- parser.add_argument('--checkpoints_dir', type=str, default='./result', help='models are save here')
- parser.add_argument('--which_iter', type=str, default='latest', help='which iterations to load')
- parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0, 1, 2 use -1 for CPU')
- parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
- parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
-
-
- # input/output sizes
- parser.add_argument('--batchSize', type=int, default=8, help='input batch size')
- parser.add_argument('--old_size', type=int, default=(256,256), help='Scale images to this size. The final image will be cropped to --crop_size.')
- parser.add_argument('--load_size', type=int, default=1024, help='Scale images to this size. The final image will be cropped to --crop_size.')
- parser.add_argument('--structure_nc', type=int, default=18 )
- parser.add_argument('--image_nc', type=int, default=3 )
-
- # for setting inputs
- parser.add_argument('--dataroot', type=str, default='./dataset/fashion/')
- parser.add_argument('--dataset_mode', type=str, default='fashion')
- parser.add_argument('--fid_gt_path', type=str)
- parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
- parser.add_argument('--nThreads', default=8, type=int, help='# threads for loading data')
- parser.add_argument('--max_dataset_size', type=int, default=sys.maxsize, help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
-
- # display parameter define
- parser.add_argument('--display_winsize', type=int, default=256, help='display window size')
- parser.add_argument('--display_id', type=int, default=1, help='display id of the web')
- parser.add_argument('--display_port', type=int, default=8096, help='visidom port of the web display')
- parser.add_argument('--display_single_pane_ncols', type=int, default=0, help='if positive, display all images in a single visidom web panel')
- parser.add_argument('--display_env', type=str, default=parser.parse_known_args()[0].name.replace('_',''), help='the environment of visidom display')
-
- return parser
-
- def gather_options(self):
- """Add additional model-specific options"""
-
- if not self.initialized:
- parser = self.initialize(self.parser)
-
- # get basic options
- opt, _ = parser.parse_known_args()
-
- # modify the options for different models
- model_option_set = model.get_option_setter(opt.model)
- parser = model_option_set(parser, self.isTrain)
-
- data_option_set = data.get_option_setter(opt.dataset_mode)
- parser = data_option_set(parser, self.isTrain)
-
-
- opt = parser.parse_args()
-
- return opt
-
- def parse(self):
- """Parse the options"""
-
- opt = self.gather_options()
- opt.isTrain = self.isTrain
-
-
- if opt.phase != 'val':
- self.print_options(opt)
-
- if torch.cuda.is_available():
- opt.device = torch.device("cuda")
- torch.backends.cudnn.benchmark = True # cudnn auto-tuner
- else:
- opt.device = torch.device("cpu")
-
- # set gpu ids
- str_ids = opt.gpu_ids.split(',')
- opt.gpu_ids = []
- for str_id in str_ids:
- id = int(str_id)
- if id >= 0:
- opt.gpu_ids.append(id)
- if len(opt.gpu_ids):
- torch.cuda.set_device(opt.gpu_ids[0])
-
- self.opt = opt
-
- return self.opt
-
- @staticmethod
- def print_options(opt):
- """print and save options"""
-
- print('--------------Options--------------')
- for k, v in sorted(vars(opt).items()):
- print('%s: %s' % (str(k), str(v)))
- print('----------------End----------------')
-
- # save to the disk
- expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
- util.mkdirs(expr_dir)
- if opt.isTrain:
- file_name = os.path.join(expr_dir, 'train_opt.txt')
- else:
- file_name = os.path.join(expr_dir, 'test_opt.txt')
- with open(file_name, 'wt') as opt_file:
- opt_file.write('--------------Options--------------\n')
- for k, v in sorted(vars(opt).items()):
- opt_file.write('%s: %s\n' % (str(k), str(v)))
- opt_file.write('----------------End----------------\n')
|