def init_training_settings(self): train_opt = self.opt['train'] # define network net_d self.net_d = define_network(deepcopy(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 = define_network(deepcopy(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) cri_gan_cls = getattr(loss_module, train_opt['gan_opt'].pop('type')) self.cri_gan = cri_gan_cls(**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 = define_network(deepcopy(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): train_opt = self.opt['train'] # define network net_d self.net_d = define_network(deepcopy(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_model_d', None) if load_path is not None: self.load_network(self.net_d, load_path, self.opt['path']['strict_load']) self.net_g.train() self.net_d.train() # define losses if train_opt.get('pixel_opt'): pixel_type = train_opt['pixel_opt'].pop('type') cri_pix_cls = getattr(loss_module, pixel_type) self.cri_pix = cri_pix_cls(**train_opt['pixel_opt']).to( self.device) else: self.cri_pix = None if train_opt.get('perceptual_opt'): percep_type = train_opt['perceptual_opt'].pop('type') cri_perceptual_cls = getattr(loss_module, percep_type) self.cri_perceptual = cri_perceptual_cls( **train_opt['perceptual_opt']).to(self.device) else: self.cri_perceptual = None if train_opt.get('gan_opt'): gan_type = train_opt['gan_opt'].pop('type') cri_gan_cls = getattr(loss_module, gan_type) self.cri_gan = cri_gan_cls(**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 = define_network(deepcopy(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_model_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']['strict_load'], 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()