File size: 5,657 Bytes
2302223 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
from argparse import ArgumentParser
from configs.paths_config import model_paths
class TrainOptions:
def __init__(self):
self.parser = ArgumentParser()
self.initialize()
def initialize(self):
self.parser.add_argument('--exp_dir', type=str,
help='Path to experiment output directory')
self.parser.add_argument('--dataset_type', default='ffhq_aging', type=str,
help='Type of dataset/experiment to run')
self.parser.add_argument('--input_nc', default=4, type=int,
help='Number of input image channels to the psp encoder')
self.parser.add_argument('--label_nc', default=0, type=int,
help='Number of input label channels to the psp encoder')
self.parser.add_argument('--output_size', default=1024, type=int,
help='Output size of generator')
self.parser.add_argument('--batch_size', default=4, type=int,
help='Batch size for training')
self.parser.add_argument('--test_batch_size', default=2, type=int,
help='Batch size for testing and inference')
self.parser.add_argument('--workers', default=4, type=int,
help='Number of train dataloader workers')
self.parser.add_argument('--test_workers', default=2, type=int,
help='Number of test/inference dataloader workers')
self.parser.add_argument('--learning_rate', default=0.0001, type=float,
help='Optimizer learning rate')
self.parser.add_argument('--optim_name', default='ranger', type=str,
help='Which optimizer to use')
self.parser.add_argument('--train_decoder', action='store_true',
help='Whether to train the decoder model')
self.parser.add_argument('--start_from_latent_avg', action='store_true',
help='Whether to add average latent vector to generate codes from encoder.')
self.parser.add_argument('--start_from_encoded_w_plus', action='store_true',
help='Whether to learn residual wrt w+ of encoded image using pretrained pSp.')
self.parser.add_argument('--lpips_lambda', default=0, type=float,
help='LPIPS loss multiplier factor')
self.parser.add_argument('--id_lambda', default=0, type=float,
help='ID loss multiplier factor')
self.parser.add_argument('--l2_lambda', default=0, type=float,
help='L2 loss multiplier factor')
self.parser.add_argument('--w_norm_lambda', default=0, type=float,
help='W-norm loss multiplier factor')
self.parser.add_argument('--aging_lambda', default=0, type=float,
help='Aging loss multiplier factor')
self.parser.add_argument('--cycle_lambda', default=0, type=float,
help='Cycle loss multiplier factor')
self.parser.add_argument('--lpips_lambda_crop', default=0, type=float,
help='LPIPS loss multiplier factor for inner image region')
self.parser.add_argument('--l2_lambda_crop', default=0, type=float,
help='L2 loss multiplier factor for inner image region')
self.parser.add_argument('--lpips_lambda_aging', default=0, type=float,
help='LPIPS loss multiplier factor for aging')
self.parser.add_argument('--l2_lambda_aging', default=0, type=float,
help='L2 loss multiplier factor for aging')
self.parser.add_argument('--stylegan_weights', default=model_paths['stylegan_ffhq'], type=str,
help='Path to StyleGAN model weights')
self.parser.add_argument('--checkpoint_path', default=None, type=str,
help='Path to pSp model checkpoint')
self.parser.add_argument('--max_steps', default=500000, type=int,
help='Maximum number of training steps')
self.parser.add_argument('--image_interval', default=100, type=int,
help='Interval for logging train images during training')
self.parser.add_argument('--board_interval', default=50, type=int,
help='Interval for logging metrics to tensorboard')
self.parser.add_argument('--val_interval', default=1000, type=int,
help='Validation interval')
self.parser.add_argument('--save_interval', default=None, type=int,
help='Model checkpoint interval')
# arguments for aging
self.parser.add_argument('--target_age', default=None, type=str,
help='Target age for training. Use `uniform_random` for random sampling of target age')
self.parser.add_argument('--use_weighted_id_loss', action="store_true",
help="Whether to weight id loss based on change in age (more change -> less weight)")
self.parser.add_argument('--pretrained_psp_path', default=model_paths['pretrained_psp'], type=str,
help="Path to pretrained pSp network.")
def parse(self):
opts = self.parser.parse_args()
return opts
|