def init_training_settings(self):
train_opt = self.opt['train']
# define network net_d
self.net_d = networks.define_net_d(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_model_d', None)
if load_path is not None:
self.load_network(self.net_d, load_path,
self.opt['path']['strict_load'])
# 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 = networks.define_net_g(deepcopy(
self.opt['network_g'])).to(self.device)
# load pretrained model
load_path = self.opt['path'].get('pretrain_model_g', None)
if load_path is not None:
self.load_network(self.net_g_ema, load_path,
self.opt['path']['strict_load'], '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 = networks.define_net_g(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_model_g', None)
if load_path is not None:
self.load_network(self.net_g, load_path,
self.opt['path']['strict_load'])
if self.is_train:
self.init_training_settings()
def __init__(self, opt):
super(SRModel, self).__init__(opt)
# define network
self.net_g = networks.define_net_g(deepcopy(opt['network_g']))
self.net_g = self.model_to_device(self.net_g)
self.print_network(self.net_g)
self.offset_frame = []
self.offset_mask = []
self.avg_feat_gt = []
self.avg_feat_out = []
# load pretrained models
load_path = self.opt['path'].get('pretrain_model_g', None)
if load_path is not None:
self.load_network(self.net_g, load_path,
self.opt['path']['strict_load'])
##load arcface
load_pth_name = '/home/wei/exp/EDVR/arcfacemodel/model_ir_se50.pth'
pretrain_dict = torch.load(load_pth_name,
map_location=lambda storage, loc: storage)
# conv1_weight_new=np.zeros( (64,5,7,7) )
# conv1_weight_new[:,:3,:,:]=pretrain_dict['conv1.weight'].cpu().data
# pretrain_dict['conv1.weight']=torch.from_numpy(conv1_weight_new )
state_dict = self.net_g.state_dict()
model_dict = state_dict
keys = state_dict.keys()
# print('state_dict',state_dict)
for k, v in pretrain_dict.items():
# kk='backbone.'+k
kk = k
for key in keys:
if kk in key:
model_dict[key] = v
break
state_dict.update(model_dict)
self.net_g.load_state_dict(state_dict)
# param.requires_grad = True
if self.is_train:
self.init_training_settings()
def __init__(self, opt):
super(UNetModel, self).__init__(opt)
# define network
self.net_g = networks.define_net_g(deepcopy(opt['network_g']))
# Define additional buffers
self.output_transform_for_loss = self.opt['output_transform_for_loss']
print('Output transform (WT_HF): {}'.format(self.output_transform_for_loss))
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_model_g', None)
if load_path is not None:
self.load_network(self.net_g, load_path,
self.opt['path']['strict_load'])
if self.is_train:
self.init_training_settings()
def __init__(self, opt):
super(StyleGAN2Model, self).__init__(opt)
# define network net_g
self.net_g = networks.define_net_g(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']
self.fixed_sample = torch.randn(16,
self.num_style_feat,
device=self.device)
if self.is_train:
self.init_training_settings()
