def init_training_settings(self): train_opt = self.opt['train'] self.ema_decay = train_opt.get('ema_decay', 0) if self.ema_decay > 0: logger = get_root_logger() logger.info( f'Use Exponential Moving Average with decay: {self.ema_decay}') # define network net_g with Exponential Moving Average (EMA) # net_g_ema is used only for testing on one GPU and saving # There is no need to wrap with DistributedDataParallel self.net_g_ema = build_network(self.opt['network_g']).to( self.device) # load pretrained model load_path = self.opt['path'].get('pretrain_network_g', None) if load_path is not None: self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema') else: self.model_ema(0) # copy net_g weight self.net_g_ema.eval() # define network net_d self.net_d = build_network(self.opt['network_d']) self.net_d = self.model_to_device(self.net_d) self.print_network(self.net_d) # load pretrained models load_path = self.opt['path'].get('pretrain_network_d', None) if load_path is not None: param_key = self.opt['path'].get('param_key_d', 'params') self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True), param_key) self.net_g.train() self.net_d.train() # define losses if train_opt.get('pixel_opt'): self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device) else: self.cri_pix = None if train_opt.get('perceptual_opt'): self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to( self.device) else: self.cri_perceptual = None if train_opt.get('gan_opt'): self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device) self.net_d_iters = train_opt.get('net_d_iters', 1) self.net_d_init_iters = train_opt.get('net_d_init_iters', 0) # set up optimizers and schedulers self.setup_optimizers() self.setup_schedulers()
def init_training_settings(self): train_opt = self.opt['train'] # define network net_d self.net_d = build_network(self.opt['network_d']) self.net_d = self.model_to_device(self.net_d) self.print_network(self.net_d) # load pretrained model load_path = self.opt['path'].get('pretrain_network_d', None) if load_path is not None: self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True)) # define network net_g with Exponential Moving Average (EMA) # net_g_ema only used for testing on one GPU and saving, do not need to # wrap with DistributedDataParallel self.net_g_ema = build_network(self.opt['network_g']).to(self.device) # load pretrained model load_path = self.opt['path'].get('pretrain_network_g', None) if load_path is not None: self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema') else: self.model_ema(0) # copy net_g weight self.net_g.train() self.net_d.train() self.net_g_ema.eval() # define losses # gan loss (wgan) self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device) # regularization weights self.r1_reg_weight = train_opt['r1_reg_weight'] # for discriminator self.path_reg_weight = train_opt['path_reg_weight'] # for generator self.net_g_reg_every = train_opt['net_g_reg_every'] self.net_d_reg_every = train_opt['net_d_reg_every'] self.mixing_prob = train_opt['mixing_prob'] self.mean_path_length = 0 # set up optimizers and schedulers self.setup_optimizers() self.setup_schedulers()
def __init__(self, opt): super(SRModel, self).__init__(opt) # define network self.net_g = build_network(opt['network_g']) self.net_g = self.model_to_device(self.net_g) self.print_network(self.net_g) # load pretrained models load_path = self.opt['path'].get('pretrain_network_g', None) if load_path is not None: self.load_network(self.net_g, load_path, self.opt['path'].get('strict_load_g', True)) if self.is_train: self.init_training_settings()
def init_training_settings(self): self.net_g.train() train_opt = self.opt["train"] self.ema_decay = train_opt.get("ema_decay", 0) if self.ema_decay > 0: logger = get_root_logger() logger.info( f"Use Exponential Moving Average with decay: {self.ema_decay}") # define network net_g with Exponential Moving Average (EMA) # net_g_ema is used only for testing on one GPU and saving # There is no need to wrap with DistributedDataParallel self.net_g_ema = build_network(self.opt["network_g"]).to( self.device) # load pretrained model load_path = self.opt["path"].get("pretrain_network_g", None) if load_path is not None: self.load_network( self.net_g_ema, load_path, self.opt["path"].get("strict_load_g", True), "params_ema", ) else: self.model_ema(0) # copy net_g weight self.net_g_ema.eval() # define losses if train_opt.get("pixel_opt"): self.cri_pix = build_loss(train_opt["pixel_opt"]).to(self.device) else: self.cri_pix = None if train_opt.get("perceptual_opt"): self.cri_perceptual = build_loss(train_opt["perceptual_opt"]).to( self.device) else: self.cri_perceptual = None if self.cri_pix is None and self.cri_perceptual is None: raise ValueError("Both pixel and perceptual losses are None.") # set up optimizers and schedulers self.setup_optimizers() self.setup_schedulers()
def init_training_settings(self): train_opt = self.opt['train'] self.ema_decay = train_opt.get('ema_decay', 0) if self.ema_decay > 0: raise (NotImplementedError( 'HiFaceGAN does not support EMA now. Pass')) self.net_g.train() self.net_d = build_network(self.opt['network_d']) self.net_d = self.model_to_device(self.net_d) self.print_network(self.net_d) # define losses # HiFaceGAN does not use pixel loss by default if train_opt.get('pixel_opt'): self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device) else: self.cri_pix = None if train_opt.get('perceptual_opt'): self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to( self.device) else: self.cri_perceptual = None if train_opt.get('feature_matching_opt'): self.cri_feat = build_loss(train_opt['feature_matching_opt']).to( self.device) else: self.cri_feat = None if self.cri_pix is None and self.cri_perceptual is None: raise ValueError('Both pixel and perceptual losses are None.') if train_opt.get('gan_opt'): self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device) self.net_d_iters = train_opt.get('net_d_iters', 1) self.net_d_init_iters = train_opt.get('net_d_init_iters', 0) # set up optimizers and schedulers self.setup_optimizers() self.setup_schedulers()
def __init__(self, opt): super(StyleGAN2Model, self).__init__(opt) # define network net_g self.net_g = build_network(opt['network_g']) self.net_g = self.model_to_device(self.net_g) self.print_network(self.net_g) # load pretrained model load_path = self.opt['path'].get('pretrain_network_g', None) if load_path is not None: param_key = self.opt['path'].get('param_key_g', 'params') self.load_network(self.net_g, load_path, self.opt['path'].get('strict_load_g', True), param_key) # latent dimension: self.num_style_feat self.num_style_feat = opt['network_g']['num_style_feat'] num_val_samples = self.opt['val'].get('num_val_samples', 16) self.fixed_sample = torch.randn(num_val_samples, self.num_style_feat, device=self.device) if self.is_train: self.init_training_settings()
def __init__(self, opt): super(SRModel, self).__init__(opt) # define network self.net_g = build_network(opt["network_g"]) self.net_g = self.model_to_device(self.net_g) self.print_network(self.net_g) # load pretrained models load_path = self.opt["path"].get("pretrain_network_g", None) if load_path is not None: param_key = self.opt["path"].get("param_key_g", "params") self.load_network( self.net_g, load_path, self.opt["path"].get("strict_load_g", True), param_key, ) if self.is_train: self.init_training_settings()
def __init__(self, opt): super(GFPGANModel, self).__init__(opt) self.idx = 0 # it is used for saving data for check # define network self.net_g = build_network(opt['network_g']) self.net_g = self.model_to_device(self.net_g) self.print_network(self.net_g) # load pretrained model load_path = self.opt['path'].get('pretrain_network_g', None) if load_path is not None: param_key = self.opt['path'].get('param_key_g', 'params') self.load_network(self.net_g, load_path, self.opt['path'].get('strict_load_g', True), param_key) self.log_size = int(math.log(self.opt['network_g']['out_size'], 2)) if self.is_train: self.init_training_settings()
def init_training_settings(self): train_opt = self.opt['train'] # define network net_d self.net_d = build_network(self.opt['network_d']) self.net_d = self.model_to_device(self.net_d) self.print_network(self.net_d) # load pretrained models load_path = self.opt['path'].get('pretrain_network_d', None) if load_path is not None: self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True)) self.net_g.train() self.net_d.train() # define losses if train_opt.get('pixel_opt'): self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device) else: self.cri_pix = None if train_opt.get('perceptual_opt'): self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to( self.device) else: self.cri_perceptual = None if train_opt.get('gan_opt'): self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device) self.net_d_iters = train_opt.get('net_d_iters', 1) self.net_d_init_iters = train_opt.get('net_d_init_iters', 0) # set up optimizers and schedulers self.setup_optimizers() self.setup_schedulers()
def init_training_settings(self): train_opt = self.opt['train'] # ----------- define net_d ----------- # self.net_d = build_network(self.opt['network_d']) self.net_d = self.model_to_device(self.net_d) self.print_network(self.net_d) # load pretrained model load_path = self.opt['path'].get('pretrain_network_d', None) if load_path is not None: self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True)) # ----------- define net_g with Exponential Moving Average (EMA) ----------- # # net_g_ema only used for testing on one GPU and saving. There is no need to wrap with DistributedDataParallel self.net_g_ema = build_network(self.opt['network_g']).to(self.device) # load pretrained model load_path = self.opt['path'].get('pretrain_network_g', None) if load_path is not None: self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema') else: self.model_ema(0) # copy net_g weight self.net_g.train() self.net_d.train() self.net_g_ema.eval() # ----------- facial component networks ----------- # if ('network_d_left_eye' in self.opt and 'network_d_right_eye' in self.opt and 'network_d_mouth' in self.opt): self.use_facial_disc = True else: self.use_facial_disc = False if self.use_facial_disc: # left eye self.net_d_left_eye = build_network(self.opt['network_d_left_eye']) self.net_d_left_eye = self.model_to_device(self.net_d_left_eye) self.print_network(self.net_d_left_eye) load_path = self.opt['path'].get('pretrain_network_d_left_eye') if load_path is not None: self.load_network(self.net_d_left_eye, load_path, True, 'params') # right eye self.net_d_right_eye = build_network( self.opt['network_d_right_eye']) self.net_d_right_eye = self.model_to_device(self.net_d_right_eye) self.print_network(self.net_d_right_eye) load_path = self.opt['path'].get('pretrain_network_d_right_eye') if load_path is not None: self.load_network(self.net_d_right_eye, load_path, True, 'params') # mouth self.net_d_mouth = build_network(self.opt['network_d_mouth']) self.net_d_mouth = self.model_to_device(self.net_d_mouth) self.print_network(self.net_d_mouth) load_path = self.opt['path'].get('pretrain_network_d_mouth') if load_path is not None: self.load_network(self.net_d_mouth, load_path, True, 'params') self.net_d_left_eye.train() self.net_d_right_eye.train() self.net_d_mouth.train() # ----------- define facial component gan loss ----------- # self.cri_component = build_loss(train_opt['gan_component_opt']).to( self.device) # ----------- define losses ----------- # # pixel loss if train_opt.get('pixel_opt'): self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device) else: self.cri_pix = None # perceptual loss if train_opt.get('perceptual_opt'): self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to( self.device) else: self.cri_perceptual = None # L1 loss is used in pyramid loss, component style loss and identity loss self.cri_l1 = build_loss(train_opt['L1_opt']).to(self.device) # gan loss (wgan) self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device) # ----------- define identity loss ----------- # if 'network_identity' in self.opt: self.use_identity = True else: self.use_identity = False if self.use_identity: # define identity network self.network_identity = build_network(self.opt['network_identity']) self.network_identity = self.model_to_device(self.network_identity) self.print_network(self.network_identity) load_path = self.opt['path'].get('pretrain_network_identity') if load_path is not None: self.load_network(self.network_identity, load_path, True, None) self.network_identity.eval() for param in self.network_identity.parameters(): param.requires_grad = False # regularization weights self.r1_reg_weight = train_opt['r1_reg_weight'] # for discriminator self.net_d_iters = train_opt.get('net_d_iters', 1) self.net_d_init_iters = train_opt.get('net_d_init_iters', 0) self.net_d_reg_every = train_opt['net_d_reg_every'] # set up optimizers and schedulers self.setup_optimizers() self.setup_schedulers()