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()