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