def get_lr_scheduler(self, optimizer, opt, mode=None):
assert mode in ['G', 'D', None]
fpm = lambda x: x + '_' + mode if mode else x
lr_scheme = opt['lr_scheme']
if lr_scheme == 'MultiStepLR_Restart':
scheduler = lr_scheduler.MultiStepLR_Restart(optimizer, opt[fpm('lr_steps')],
restarts=opt[fpm('restarts')], weights=opt[fpm('restart_weights')],
gamma=opt[fpm('lr_gamma')], clear_state=opt[fpm('clear_state')])
elif lr_scheme == 'CosineAnnealingLR_Restart':
scheduler = lr_scheduler.CosineAnnealingLR_Restart(optimizer, opt[fpm('T_period')],
eta_min=opt[fpm('eta_min')], restarts=opt[fpm('restarts')],
weights=opt[fpm('restart_weights')])
elif lr_scheme == 'StepLR':
scheduler = lr_scheduler.StepLR(optimizer,
step_size=opt[fpm('lr_step')], gamma=opt[fpm('lr_gamma')])
elif lr_scheme == 'MultiStepLR':
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=opt[fpm('lr_steps')], gamma=opt[fpm('lr_gamma')])
elif lr_scheme == 'CosineAnnealingLR':
scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
T_max=opt[fpm('T_max')], eta_min=opt[fpm('eta_min')])
else:
raise TypeError('Unknown lr scheduler type: {}'.format(lr_scheme))
return scheduler
def __init__(self, opt):
super(P_Model, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
# define network and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
# print network
self.print_network()
self.load()
if self.is_train:
#self.init_model()
self.netG.train()
# loss
loss_type = train_opt['pixel_criterion']
if loss_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif loss_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
elif loss_type == 'cb':
self.cri_pix = CharbonnierLoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] is not recognized.'.format(loss_type))
self.l_pix_w = train_opt['pixel_weight']
# optimizers
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1'],
train_opt['beta2']))
#self.optimizer_G = torch.optim.SGD(optim_params, lr=train_opt['lr_G'], momentum=0.9)
self.optimizers.append(self.optimizer_G)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
print('MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
def __init__(self, opt):
super(SRModel, self).__init__(opt)
train_opt = opt['train']
if self.is_train:
self.netG.train()
# loss
loss_type = train_opt['pixel_criterion']
if loss_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif loss_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] is not recognized.'.format(loss_type))
self.l_pix_w = train_opt['pixel_weight']
# optimizers
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
if opt["network_G"]["which_model_G"] == "fsrcnn":
self.optimizer_G = torch.optim.Adam(
[{
'params': self.netG.first_part.parameters()
}, {
'params': self.netG.mid_part.parameters()
}, {
'params': self.netG.last_part.parameters(),
'lr': train_opt['lr_G'] * 0.1
}],
lr=train_opt['lr_G'])
else:
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G)
self.optimizers.append(self.optimizer_G)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'Choose MultiStepLR or CosineAnnealingLR')
self.log_dict = OrderedDict()
def init_optimizer_and_scheduler(self, train_opt):
# optimizers
self.optimizers = []
wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
optim_params_RRDB = []
optim_params_other = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
print(k, v.requires_grad)
if v.requires_grad:
if '.RRDB.' in k:
optim_params_RRDB.append(v)
print('opt', k)
else:
optim_params_other.append(v)
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
print('rrdb params', len(optim_params_RRDB))
self.optimizer_G = torch.optim.Adam(
[{
"params": optim_params_other,
"lr": train_opt['lr_G'],
'beta1': train_opt['beta1'],
'beta2': train_opt['beta2'],
'weight_decay': wd_G
}, {
"params": optim_params_RRDB,
"lr": train_opt.get('lr_RRDB', train_opt['lr_G']),
'beta1': train_opt['beta1'],
'beta2': train_opt['beta2'],
'weight_decay': wd_G
}], )
self.optimizers.append(self.optimizer_G)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state'],
lr_steps_invese=train_opt.get('lr_steps_inverse', [])))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
def __init__(self, opt):
super(SRGANModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
self.train_opt = train_opt
self.opt = opt
self.segmentor = None
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
if self.is_train:
self.netD = networks.define_D(opt).to(self.device)
if train_opt.get("gan_video_weight", 0) > 0:
self.net_video_D = networks.define_video_D(opt).to(self.device)
if opt['dist']:
self.netD = DistributedDataParallel(
self.netD, device_ids=[torch.cuda.current_device()])
if train_opt.get("gan_video_weight", 0) > 0:
self.net_video_D = DistributedDataParallel(
self.net_video_D,
device_ids=[torch.cuda.current_device()])
else:
self.netD = DataParallel(self.netD)
if train_opt.get("gan_video_weight", 0) > 0:
self.net_video_D = DataParallel(self.net_video_D)
self.netG.train()
self.netD.train()
if train_opt.get("gan_video_weight", 0) > 0:
self.net_video_D.train()
# define losses, optimizer and scheduler
if self.is_train:
# G pixel loss
if train_opt['pixel_weight'] > 0:
l_pix_type = train_opt['pixel_criterion']
if l_pix_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif l_pix_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_pix_type))
self.l_pix_w = train_opt['pixel_weight']
else:
logger.info('Remove pixel loss.')
self.cri_pix = None
# Pixel mask loss
if train_opt.get("pixel_mask_weight", 0) > 0:
l_pix_type = train_opt['pixel_mask_criterion']
self.cri_pix_mask = LMaskLoss(
l_pix_type=l_pix_type,
segm_mask=train_opt['segm_mask']).to(self.device)
self.l_pix_mask_w = train_opt['pixel_mask_weight']
else:
logger.info('Remove pixel mask loss.')
self.cri_pix_mask = None
# G feature loss
if train_opt['feature_weight'] > 0:
l_fea_type = train_opt['feature_criterion']
if l_fea_type == 'l1':
self.cri_fea = nn.L1Loss().to(self.device)
elif l_fea_type == 'l2':
self.cri_fea = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_fea_type))
self.l_fea_w = train_opt['feature_weight']
else:
logger.info('Remove feature loss.')
self.cri_fea = None
if self.cri_fea: # load VGG perceptual loss
self.netF = networks.define_F(opt,
use_bn=False).to(self.device)
if opt['dist']:
self.netF = DistributedDataParallel(
self.netF, device_ids=[torch.cuda.current_device()])
else:
self.netF = DataParallel(self.netF)
# GD gan loss
self.cri_gan = GANLoss(train_opt['gan_type'], 1.0,
0.0).to(self.device)
self.l_gan_w = train_opt['gan_weight']
# Video gan weight
if train_opt.get("gan_video_weight", 0) > 0:
self.cri_video_gan = GANLoss(train_opt['gan_video_type'], 1.0,
0.0).to(self.device)
self.l_gan_video_w = train_opt['gan_video_weight']
# can't use optical flow with i and i+1 because we need i+2 lr to calculate i+1 oflow
if 'train' in self.opt['datasets'].keys():
key = "train"
else:
key = 'test_1'
assert self.opt['datasets'][key][
'optical_flow_with_ref'] == True, f"Current value = {self.opt['datasets'][key]['optical_flow_with_ref']}"
# D_update_ratio and D_init_iters
self.D_update_ratio = train_opt['D_update_ratio'] if train_opt[
'D_update_ratio'] else 1
self.D_init_iters = train_opt['D_init_iters'] if train_opt[
'D_init_iters'] else 0
# optimizers
# G
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1_G'],
train_opt['beta2_G']))
self.optimizers.append(self.optimizer_G)
# D
wd_D = train_opt['weight_decay_D'] if train_opt[
'weight_decay_D'] else 0
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=train_opt['lr_D'],
weight_decay=wd_D,
betas=(train_opt['beta1_D'],
train_opt['beta2_D']))
self.optimizers.append(self.optimizer_D)
# Video D
if train_opt.get("gan_video_weight", 0) > 0:
self.optimizer_video_D = torch.optim.Adam(
self.net_video_D.parameters(),
lr=train_opt['lr_D'],
weight_decay=wd_D,
betas=(train_opt['beta1_D'], train_opt['beta2_D']))
self.optimizers.append(self.optimizer_video_D)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
self.print_network() # print network
self.load() # load G and D if needed
def __init__(self, opt: dict):
super(IRNModel, self).__init__(opt)
self.train_opt: dict = opt["train"]
# self.test_opt: str = opt['test']
# TODO : fix
self.netG = networks.define_G(opt).to(self.device)
# self.netG = networks.define_G(opt)
self.netG = DataParallel(self.netG)
self.print_network()
# self.load() TODO
# TODO
self.Quantization = Quantization()
if self.is_train:
self.netG.train()
self.Reconstruction_forw = ReconstructionLoss(
losstype=self.train_opt["pixel_criterion_forw"])
self.Reconstruction_back = ReconstructionLoss(
losstype=self.train_opt["pixel_criterion_back"])
wd_G = (self.train_opt["weight_decay_G"]
if self.train_opt["weight_decay_G"] else 0)
optim_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
optim_params.append(v)
else:
logger.warning(
"Params [{:s}] will not optimize.".format(k))
self.optimizer_G = torch.optim.Adam(
optim_params,
lr=self.train_opt["lr_G"],
weight_decay=wd_G,
betas=(self.train_opt["beta1"], self.train_opt["beta2"]),
)
self.optimizers.append(self.optimizer_G)
if self.train_opt["lr_scheme"] == "MultiStepLR":
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
self.train_opt["lr_steps"],
restarts=self.train_opt["restarts"],
weights=self.train_opt["restart_weights"],
gamma=self.train_opt["lr_gamma"],
clear_state=self.train_opt["clear_state"],
))
elif self.train_opt["lr_scheme"] == "CosineAnnealingLR_Restart":
pass
# TODO
else:
raise NotImplementedError(
"MultiStepLR learning rate scheme is enough.")
self.log_dict: dict = OrderedDict()
def __init__(self, opt):
super(RRSNetModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
self.l1_init = train_opt['l1_init'] if train_opt['l1_init'] else 0
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG,
device_ids=[torch.cuda.current_device()],
find_unused_parameters=True)
else:
self.netG = DataParallel(self.netG)
if self.is_train:
self.netG.train()
self.load() # load G and D if needed
# define losses, optimizer and scheduler
if self.is_train:
# G pixel loss
if train_opt['pixel_weight'] > 0:
l_pix_type = train_opt['pixel_criterion']
if l_pix_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif l_pix_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_pix_type))
self.l_pix_w = train_opt['pixel_weight']
else:
logger.info('Remove pixel loss.')
self.cri_pix = None
# Branch_init_iters
self.Branch_pretrain = train_opt['Branch_pretrain'] if train_opt[
'Branch_pretrain'] else 0
self.Branch_init_iters = train_opt[
'Branch_init_iters'] if train_opt['Branch_init_iters'] else 1
# gradient_pixel_loss
self.cri_pix_grad = nn.MSELoss().to(self.device)
self.l_pix_grad_w = train_opt['gradient_pixel_weight']
# optimizers
# G
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1_G'],
train_opt['beta2_G']))
self.optimizers.append(self.optimizer_G)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
self.get_grad = Get_gradient()
self.get_grad_nopadding = Get_gradient_nopadding()
self.print_network() # print network
def __init__(self, opt):
super(IRNpModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
test_opt = opt['test']
self.train_opt = train_opt
self.test_opt = test_opt
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
# print network
self.print_network()
self.load()
self.Quantization = Quantization()
if self.is_train:
self.netD = networks.define_D(opt).to(self.device)
if opt['dist']:
self.netD = DistributedDataParallel(
self.netD, device_ids=[torch.cuda.current_device()])
else:
self.netD = DataParallel(self.netD)
self.netG.train()
self.netD.train()
# loss
self.Reconstruction_forw = ReconstructionLoss(
losstype=self.train_opt['pixel_criterion_forw'])
self.Reconstruction_back = ReconstructionLoss(
losstype=self.train_opt['pixel_criterion_back'])
# feature loss
if train_opt['feature_weight'] > 0:
self.Reconstructionf = ReconstructionLoss(
losstype=self.train_opt['feature_criterion'])
self.l_fea_w = train_opt['feature_weight']
self.netF = networks.define_F(opt,
use_bn=False).to(self.device)
if opt['dist']:
self.netF = DistributedDataParallel(
self.netF, device_ids=[torch.cuda.current_device()])
else:
self.netF = DataParallel(self.netF)
else:
self.l_fea_w = 0
# GD gan loss
self.cri_gan = GANLoss(train_opt['gan_type'], 1.0,
0.0).to(self.device)
self.l_gan_w = train_opt['gan_weight']
# D_update_ratio and D_init_iters
self.D_update_ratio = train_opt['D_update_ratio'] if train_opt[
'D_update_ratio'] else 1
self.D_init_iters = train_opt['D_init_iters'] if train_opt[
'D_init_iters'] else 0
# optimizers
# G
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1'],
train_opt['beta2']))
self.optimizers.append(self.optimizer_G)
# D
wd_D = train_opt['weight_decay_D'] if train_opt[
'weight_decay_D'] else 0
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=train_opt['lr_D'],
weight_decay=wd_D,
betas=(train_opt['beta1_D'],
train_opt['beta2_D']))
self.optimizers.append(self.optimizer_D)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
def __init__(self, opt):
super(RRDBSRModel, self).__init__(opt)
# define networks and load pretrained models
train_opt = opt['train']
self.netG_SR = networks.define_SR(opt).to(self.device)
if self.is_train:
if not self.opt['full_sr']:
self.netG_BA = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG_SR = DistributedDataParallel(
self.netG_SR,
device_ids=[torch.cuda.current_device()],
find_unused_parameters=True)
if self.is_train:
if not self.opt['full_sr']:
self.netG_BA = DistributedDataParallel(
self.netG_BA,
device_ids=[torch.cuda.current_device()],
find_unused_parameters=True)
else:
self.netG_SR = DataParallel(self.netG_SR)
if self.is_train:
if not self.opt['full_sr']:
self.netG_BA = DataParallel(self.netG_BA)
# define losses, optimizer and scheduler
if self.is_train:
# losses
self.criterion = SRLoss(train_opt['loss_type']).to(
self.device) # define GAN loss.
# optimizers
self.optimizer_G = torch.optim.Adam(self.netG_SR.parameters(),
lr=train_opt['lr'],
betas=(train_opt['beta1'],
train_opt['beta2']))
self.optimizers.append(self.optimizer_G)
#scheduler
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
else:
raise NotImplementedError("lr_scheme does not implement still")
self.log_dict = OrderedDict()
self.load() # load pre-trained mode
if self.is_train:
self.train_state()
def __init__(self, opt, is_train):
super(SRGANModel, self).__init__(opt, is_train)
train_opt = opt
self.rank = 0
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
self.netG = DataParallel(self.netG)
if self.is_train:
self.netD = networks.define_D(opt).to(self.device)
self.netD = DataParallel(self.netD)
self.netG.train()
self.netD.train()
# define losses, optimizer and scheduler
if self.is_train:
# G pixel loss
if train_opt.pixel_weight > 0:
l_pix_type = train_opt.pixel_criterion
if l_pix_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif l_pix_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_pix_type))
self.l_pix_w = train_opt.pixel_weight
else:
logger.info('Remove pixel loss.')
self.cri_pix = None
# G feature loss
if train_opt.feature_weight > 0:
l_fea_type = train_opt.feature_criterion
if l_fea_type == 'l1':
self.cri_fea = nn.L1Loss().to(self.device)
elif l_fea_type == 'l2':
self.cri_fea = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_fea_type))
self.l_fea_w = train_opt.feature_weight
else:
logger.info('Remove feature loss.')
self.cri_fea = None
if self.cri_fea: # load VGG perceptual loss
self.netF = networks.define_F(opt,
use_bn=False).to(self.device)
self.netF = DataParallel(self.netF)
# GD gan loss
self.cri_gan = GANLoss(train_opt.gan_type, 1.0,
0.0).to(self.device)
self.l_gan_w = train_opt.gan_weight
# D_update_ratio and D_init_iters
self.D_update_ratio = train_opt.D_update_ratio if train_opt.D_update_ratio else 1
self.D_init_iters = train_opt.D_init_iters if train_opt.D_init_iters else 0
# optimizers
# G
wd_G = train_opt.weight_decay_G if train_opt.weight_decay_G else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt.lr_G,
weight_decay=wd_G,
betas=(train_opt.beta1_G,
train_opt.beta2_G))
self.optimizers.append(self.optimizer_G)
# D
wd_D = train_opt.weight_decay_D if train_opt.weight_decay_D else 0
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=train_opt.lr_D,
weight_decay=wd_D,
betas=(train_opt.beta1_D,
train_opt.beta2_D))
self.optimizers.append(self.optimizer_D)
# schedulers
if train_opt.lr_scheme == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt.lr_steps,
restarts=None,
weights=None,
gamma=train_opt.lr_gamma,
clear_state=False))
elif train_opt.lr_scheme == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt.T_period,
eta_min=train_opt.eta_min,
restarts=train_opt.restarts,
weights=train_opt.restart_weights))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
self.print_network() # print network
self.load() # load G and D if needed
def __init__(self, opt):
super(SRGANModel, self).__init__(opt)
if opt["dist"]:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt["train"]
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt["dist"]:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
if self.is_train:
self.netD = networks.define_D(opt).to(self.device)
if opt["dist"]:
self.netD = DistributedDataParallel(
self.netD, device_ids=[torch.cuda.current_device()])
else:
self.netD = DataParallel(self.netD)
self.netG.train()
self.netD.train()
# define losses, optimizer and scheduler
if self.is_train:
# G pixel loss
if train_opt["pixel_weight"] > 0:
l_pix_type = train_opt["pixel_criterion"]
if l_pix_type == "l1":
self.cri_pix = nn.L1Loss().to(self.device)
elif l_pix_type == "l2":
self.cri_pix = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
"Loss type [{:s}] not recognized.".format(l_pix_type))
self.l_pix_w = train_opt["pixel_weight"]
else:
logger.info("Remove pixel loss.")
self.cri_pix = None
# G feature loss
if train_opt["feature_weight"] > 0:
l_fea_type = train_opt["feature_criterion"]
if l_fea_type == "l1":
self.cri_fea = nn.L1Loss().to(self.device)
elif l_fea_type == "l2":
self.cri_fea = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
"Loss type [{:s}] not recognized.".format(l_fea_type))
self.l_fea_w = train_opt["feature_weight"]
else:
logger.info("Remove feature loss.")
self.cri_fea = None
if self.cri_fea: # load VGG perceptual loss
self.netF = networks.define_F(opt,
use_bn=False).to(self.device)
if opt["dist"]:
self.netF = DistributedDataParallel(
self.netF, device_ids=[torch.cuda.current_device()])
else:
self.netF = DataParallel(self.netF)
# GD gan loss
self.cri_gan = GANLoss(train_opt["gan_type"], 1.0,
0.0).to(self.device)
self.l_gan_w = train_opt["gan_weight"]
# D_update_ratio and D_init_iters
self.D_update_ratio = (train_opt["D_update_ratio"]
if train_opt["D_update_ratio"] else 1)
self.D_init_iters = (train_opt["D_init_iters"]
if train_opt["D_init_iters"] else 0)
# optimizers
# G
wd_G = train_opt["weight_decay_G"] if train_opt[
"weight_decay_G"] else 0
optim_params = []
for (
k,
v,
) in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
"Params [{:s}] will not optimize.".format(k))
self.optimizer_G = torch.optim.Adam(
optim_params,
lr=train_opt["lr_G"],
weight_decay=wd_G,
betas=(train_opt["beta1_G"], train_opt["beta2_G"]),
)
self.optimizers.append(self.optimizer_G)
# D
wd_D = train_opt["weight_decay_D"] if train_opt[
"weight_decay_D"] else 0
self.optimizer_D = torch.optim.Adam(
self.netD.parameters(),
lr=train_opt["lr_D"],
weight_decay=wd_D,
betas=(train_opt["beta1_D"], train_opt["beta2_D"]),
)
self.optimizers.append(self.optimizer_D)
# schedulers
if train_opt["lr_scheme"] == "MultiStepLR":
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt["lr_steps"],
restarts=train_opt["restarts"],
weights=train_opt["restart_weights"],
gamma=train_opt["lr_gamma"],
clear_state=train_opt["clear_state"],
))
elif train_opt["lr_scheme"] == "CosineAnnealingLR_Restart":
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt["T_period"],
eta_min=train_opt["eta_min"],
restarts=train_opt["restarts"],
weights=train_opt["restart_weights"],
))
else:
raise NotImplementedError(
"MultiStepLR learning rate scheme is enough.")
self.log_dict = OrderedDict()
self.print_network() # print network
self.load() # load G and D if needed
def __init__(self, opt):
super(VideoBaseModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
# define network and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
# print network
# self.print_network()
self.load()
'''
# Freeze front when indicated
if opt['train']['freeze_front']:
print('Freeze the front part of the model')
if opt['network_G']['which_model_G'] == 'EDVR':
for name, child in self.netG.module.named_children():
if name not in ('tsa_fusion', 'recon_trunk', 'upconv1', 'upconv2', 'HRconv', 'conv_last'):
for params in child.parameters():
params.requires_grad = False
elif opt['network_G']['which_model_G'] == 'DUF':
for name, child in self.netG.module.named_children():
if name in ('conv3d_1', 'dense_block_1', 'dense_block_2', 'dense_block_3'):
for params in child.parameters():
params.requires_grad = False
else:
raise NotImplementedError()
'''
if self.is_train:
self.netG.train()
#### loss
loss_type = train_opt['pixel_criterion']
if loss_type == 'l1':
self.cri_pix = nn.L1Loss(reduction='mean').to(
self.device) # Change from sum to mean
elif loss_type == 'l2':
self.cri_pix = nn.MSELoss(reduction='mean').to(
self.device) # Change from sum to mean
elif loss_type == 'cb':
self.cri_pix = CharbonnierLoss().to(self.device)
elif loss_type == 'huber':
self.cri_pix = HuberLoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] is not recognized.'.format(loss_type))
self.l_pix_w = train_opt['pixel_weight']
#### optimizers
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
if train_opt['ft_tsa_only']:
normal_params = []
tsa_fusion_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
if 'tsa_fusion' in k:
tsa_fusion_params.append(v)
else:
normal_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
optim_params = [
{ # add normal params first
'params': normal_params,
'lr': train_opt['lr_G']
},
{
'params': tsa_fusion_params,
'lr': train_opt['lr_G']
},
]
if opt['train']['freeze_front']:
normal_params = []
freeze_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
if 'module.conv3d_1' in k or 'module.dense_block_1' in k or 'module.dense_block_2' in k or 'module.dense_block_3' in k:
freeze_params.append(v)
else:
normal_params.append(v)
optim_params = [
{ # add normal params first
'params': normal_params,
'lr': train_opt['lr_G']
},
{
'params': freeze_params,
'lr': 0
},
]
elif train_opt['small_offset_lr']:
normal_params = []
conv_offset_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
if 'pcd_align' in k or 'fea_L' in k or 'feature_extraction' in k or 'conv_first' in k:
conv_offset_params.append(v)
else:
normal_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
optim_params = [
{ # add normal params first
'params': normal_params,
'lr': train_opt['lr_G']
},
{
'params': conv_offset_params,
'lr': train_opt['lr_G'] * 0.1
},
]
else:
optim_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
if train_opt['optim'] == 'SGD':
self.optimizer_G = torch.optim.SGD(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G)
else:
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1'],
train_opt['beta2']))
self.optimizers.append(self.optimizer_G)
#### schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError()
self.log_dict = OrderedDict()
def __init__(self, opt):
super(ClassSR_Model, self).__init__(opt)
self.patch_size = int(opt["patch_size"])
self.step = int(opt["step"])
self.scale = int(opt["scale"])
self.name = opt['name']
self.which_model = opt['network_G']['which_model_G']
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
# define network and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
# print network
self.print_network()
self.load()
if self.is_train:
self.l1w = float(opt["train"]["l1w"])
self.class_loss_w = float(opt["train"]["class_loss_w"])
self.average_loss_w = float(opt["train"]["average_loss_w"])
self.pf = opt['logger']['print_freq']
self.batch_size = int(opt['datasets']['train']['batch_size'])
self.netG.train()
# loss
loss_type = train_opt['pixel_criterion']
if loss_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif loss_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
elif loss_type == 'cb':
self.cri_pix = CharbonnierLoss().to(self.device)
elif loss_type == 'ClassSR_loss':
self.cri_pix = nn.L1Loss().to(self.device)
self.class_loss = class_loss_3class().to(self.device)
self.average_loss = average_loss_3class().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] is not recognized.'.format(loss_type))
# optimizers
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
if opt['fix_SR_module']:
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad and "class" not in k:
v.requires_grad = False
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1'],
train_opt['beta2']))
self.optimizers.append(self.optimizer_G)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
def __init__(self, opt, dataset=None):
super(SRGANModel, self).__init__(opt)
if dataset:
self.cri_text = True
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
if self.is_train:
self.netD = networks.define_D(opt).to(self.device)
if opt['dist']:
self.netD = DistributedDataParallel(
self.netD, device_ids=[torch.cuda.current_device()])
else:
self.netD = DataParallel(self.netD)
self.netG.train()
self.netD.train()
# define losses, optimizer and scheduler
if self.is_train:
# G pixel loss
if train_opt['pixel_weight'] > 0:
l_pix_type = train_opt['pixel_criterion']
if l_pix_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif l_pix_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_pix_type))
self.l_pix_w = train_opt['pixel_weight']
else:
logger.info('Remove pixel loss.')
self.cri_pix = None
# G feature loss
if train_opt['feature_weight'] > 0:
l_fea_type = train_opt['feature_criterion']
if l_fea_type == 'l1':
self.cri_fea = nn.L1Loss().to(self.device)
elif l_fea_type == 'l2':
self.cri_fea = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_fea_type))
self.l_fea_w = train_opt['feature_weight']
else:
logger.info('Remove feature loss.')
self.cri_fea = None
if self.cri_fea: # load VGG perceptual loss
self.netF = networks.define_F(opt,
use_bn=False).to(self.device)
if opt['dist']:
pass # do not need to use DistributedDataParallel for netF
else:
self.netF = DataParallel(self.netF)
if self.cri_text:
from lib.models.model_builder import ModelBuilder
self.netT = ModelBuilder(
arch="ResNet_ASTER",
rec_num_classes=dataset.rec_num_classes,
sDim=512,
attDim=512,
max_len_labels=100,
eos=dataset.char2id[dataset.EOS],
STN_ON=True).to(self.device)
self.netT = DataParallel(self.netT)
self.netT.eval()
from lib.util.serialization import load_checkpoint
checkpoint = load_checkpoint(train_opt['text_model'])
self.netT.load_state_dict(checkpoint['state_dict'])
# GD gan loss
self.cri_gan = GANLoss(train_opt['gan_type'], 1.0,
0.0).to(self.device)
self.l_gan_w = train_opt['gan_weight']
# D_update_ratio and D_init_iters
self.D_update_ratio = train_opt['D_update_ratio'] if train_opt[
'D_update_ratio'] else 1
self.D_init_iters = train_opt['D_init_iters'] if train_opt[
'D_init_iters'] else 0
# optimizers
# G
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1_G'],
train_opt['beta2_G']))
self.optimizers.append(self.optimizer_G)
# D
wd_D = train_opt['weight_decay_D'] if train_opt[
'weight_decay_D'] else 0
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=train_opt['lr_D'],
weight_decay=wd_D,
betas=(train_opt['beta1_D'],
train_opt['beta2_D']))
self.optimizers.append(self.optimizer_D)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
self.print_network() # print network
self.load() # load G and D if needed
def __init__(self, opt):
super(SRGANModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
if self.is_train:
self.netD = networks.define_D(opt).to(self.device)
if opt['dist']:
self.netD = DistributedDataParallel(
self.netD, device_ids=[torch.cuda.current_device()])
else:
self.netD = DataParallel(self.netD)
self.netG.train()
self.netD.train()
# define losses, optimizer and scheduler
if self.is_train:
# G pixel loss
if train_opt['pixel_weight'] > 0:
l_pix_type = train_opt['pixel_criterion']
if l_pix_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif l_pix_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_pix_type))
self.l_pix_w = train_opt['pixel_weight']
else:
logger.info('Remove pixel loss.')
self.cri_pix = None
# G feature loss
if train_opt['feature_weight'] > 0:
l_fea_type = train_opt['feature_criterion']
if l_fea_type == 'l1':
self.cri_fea = nn.L1Loss().to(self.device)
elif l_fea_type == 'l2':
self.cri_fea = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_fea_type))
self.l_fea_w = train_opt['feature_weight']
else:
logger.info('Remove feature loss.')
self.cri_fea = None
if self.cri_fea: # load VGG perceptual loss
self.netF = networks.define_F(opt,
use_bn=False).to(self.device)
if opt['dist']:
self.netF = DistributedDataParallel(
self.netF, device_ids=[torch.cuda.current_device()])
else:
self.netF = DataParallel(self.netF)
# G Rank-content loss
if train_opt['R_weight'] > 0:
self.l_R_w = train_opt['R_weight'] # load rank-content loss
self.R_bias = train_opt['R_bias']
self.netR = networks.define_R(opt).to(self.device)
if opt['dist']:
self.netR = DistributedDataParallel(
self.netR, device_ids=[torch.cuda.current_device()])
else:
self.netR = DataParallel(self.netR)
else:
logger.info('Remove rank-content loss.')
# GD gan loss
self.cri_gan = GANLoss(train_opt['gan_type'], 1.0,
0.0).to(self.device)
self.l_gan_w = train_opt['gan_weight']
# D_update_ratio and D_init_iters
self.D_update_ratio = train_opt['D_update_ratio'] if train_opt[
'D_update_ratio'] else 1
self.D_init_iters = train_opt['D_init_iters'] if train_opt[
'D_init_iters'] else 0
# optimizers
# G
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1_G'],
train_opt['beta2_G']))
self.optimizers.append(self.optimizer_G)
# D
wd_D = train_opt['weight_decay_D'] if train_opt[
'weight_decay_D'] else 0
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=train_opt['lr_D'],
weight_decay=wd_D,
betas=(train_opt['beta1_D'],
train_opt['beta2_D']))
self.optimizers.append(self.optimizer_D)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
self.print_network() # print network
self.load() # load G and D if needed
def __init__(self, opt):
super(SRVmafModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
self.use_gpu = opt['network_G']['use_gpu']
self.use_gpu = True
self.real_IQA_only = train_opt['IQA_only']
# define network and load pretrained models
if self.use_gpu:
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
else:
self.netG = networks.define_G(opt)
if self.is_train:
if train_opt['IQA_weight']:
if train_opt['IQA_criterion'] == 'vmaf':
self.cri_IQA = nn.MSELoss()
self.l_IQA_w = train_opt['IQA_weight']
self.netI = networks.define_I(opt)
if opt['dist']:
pass
else:
self.netI = DataParallel(self.netI)
else:
logger.info('Remove IQA loss.')
self.cri_IQA = None
# print network
self.print_network()
self.load()
if self.is_train:
self.netG.train()
# pixel loss
loss_type = train_opt['pixel_criterion']
if loss_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif loss_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
elif loss_type == 'cb':
self.cri_pix = CharbonnierLoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] is not recognized.'.format(loss_type))
self.l_pix_w = train_opt['pixel_weight']
# CX loss
if train_opt['CX_weight']:
l_CX_type = train_opt['CX_criterion']
if l_CX_type == 'contextual_loss':
self.cri_CX = ContextualLoss()
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_CX_type))
self.l_CX_w = train_opt['CX_weight']
else:
logger.info('Remove CX loss.')
self.cri_CX = None
# ssim loss
if train_opt['ssim_weight']:
self.cri_ssim = train_opt['ssim_criterion']
self.l_ssim_w = train_opt['ssim_weight']
self.ssim_window = train_opt['ssim_window']
else:
logger.info('Remove ssim loss.')
self.cri_ssim = None
# load VGG perceptual loss if use CX loss
# if train_opt['CX_weight']:
# self.netF = networks.define_F(opt, use_bn=False).to(self.device)
# if opt['dist']:
# pass # do not need to use DistributedDataParallel for netF
# else:
# self.netF = DataParallel(self.netF)
# optimizers of netG
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1'],
train_opt['beta2']))
self.optimizers.append(self.optimizer_G)
# optimizers of netI
if train_opt['IQA_weight']:
wd_I = train_opt['weight_decay_I'] if train_opt[
'weight_decay_I'] else 0
optim_params = []
for k, v in self.netI.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_I = torch.optim.Adam(optim_params,
lr=train_opt['lr_I'],
weight_decay=wd_I,
betas=(train_opt['beta1'],
train_opt['beta2']))
self.optimizers.append(self.optimizer_I)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
self.set_requires_grad(self.netG, False)
self.set_requires_grad(self.netI, False)
def __init__(self, opt):
super(VideoSRBaseModel, self).__init__(opt)
if opt["dist"]:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt["train"]
# define network and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt["dist"]:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
# print network
self.print_network()
self.load()
if self.is_train:
self.netG.train()
#### loss
loss_type = train_opt["pixel_criterion"]
if loss_type == "l1":
self.cri_pix = nn.L1Loss(reduction="sum").to(self.device)
elif loss_type == "l2":
self.cri_pix = nn.MSELoss(reduction="sum").to(self.device)
elif loss_type == "cb":
self.cri_pix = CharbonnierLoss().to(self.device)
else:
raise NotImplementedError(
"Loss type [{:s}] is not recognized.".format(loss_type))
self.l_pix_w = train_opt["pixel_weight"]
#### optimizers
wd_G = train_opt["weight_decay_G"] if train_opt[
"weight_decay_G"] else 0
if train_opt["ft_tsa_only"]:
normal_params = []
tsa_fusion_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
if "tsa_fusion" in k:
tsa_fusion_params.append(v)
else:
normal_params.append(v)
else:
if self.rank <= 0:
logger.warning(
"Params [{:s}] will not optimize.".format(k))
optim_params = [
{ # add normal params first
"params": normal_params,
"lr": train_opt["lr_G"],
},
{"params": tsa_fusion_params, "lr": train_opt["lr_G"]},
]
else:
optim_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
"Params [{:s}] will not optimize.".format(k))
self.optimizer_G = torch.optim.Adam(
optim_params,
lr=train_opt["lr_G"],
weight_decay=wd_G,
betas=(train_opt["beta1"], train_opt["beta2"]),
)
self.optimizers.append(self.optimizer_G)
#### schedulers
if train_opt["lr_scheme"] == "MultiStepLR":
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt["lr_steps"],
restarts=train_opt["restarts"],
weights=train_opt["restart_weights"],
gamma=train_opt["lr_gamma"],
clear_state=train_opt["clear_state"],
))
elif train_opt["lr_scheme"] == "CosineAnnealingLR_Restart":
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt["T_period"],
eta_min=train_opt["eta_min"],
restarts=train_opt["restarts"],
weights=train_opt["restart_weights"],
))
else:
raise NotImplementedError()
self.log_dict = OrderedDict()
def __init__(self, opt):
super(B_Model, self).__init__(opt)
if opt["dist"]:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
# define network and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt["dist"]:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()]
)
else:
self.netG = DataParallel(self.netG)
# print network
self.print_network()
self.load()
if self.is_train:
train_opt = opt["train"]
# self.init_model() # Not use init is OK, since Pytorch has its owen init (by default)
self.netG.train()
# loss
loss_type = train_opt["pixel_criterion"]
if loss_type == "l1":
self.cri_pix = nn.L1Loss().to(self.device)
elif loss_type == "l2":
self.cri_pix = nn.MSELoss().to(self.device)
elif loss_type == "cb":
self.cri_pix = CharbonnierLoss().to(self.device)
else:
raise NotImplementedError(
"Loss type [{:s}] is not recognized.".format(loss_type)
)
self.l_pix_w = train_opt["pixel_weight"]
# optimizers
wd_G = train_opt["weight_decay_G"] if train_opt["weight_decay_G"] else 0
optim_params = []
for (
k,
v,
) in self.netG.named_parameters(): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning("Params [{:s}] will not optimize.".format(k))
self.optimizer_G = torch.optim.Adam(
optim_params,
lr=train_opt["lr_G"],
weight_decay=wd_G,
betas=(train_opt["beta1"], train_opt["beta2"]),
)
# self.optimizer_G = torch.optim.SGD(optim_params, lr=train_opt['lr_G'], momentum=0.9)
self.optimizers.append(self.optimizer_G)
# schedulers
if train_opt["lr_scheme"] == "MultiStepLR":
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt["lr_steps"],
restarts=train_opt["restarts"],
weights=train_opt["restart_weights"],
gamma=train_opt["lr_gamma"],
clear_state=train_opt["clear_state"],
)
)
elif train_opt["lr_scheme"] == "CosineAnnealingLR_Restart":
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt["T_period"],
eta_min=train_opt["eta_min"],
restarts=train_opt["restarts"],
weights=train_opt["restart_weights"],
)
)
else:
print("MultiStepLR learning rate scheme is enough.")
self.log_dict = OrderedDict()
def __init__(self, opt):
super(MWGANModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
self.train_opt = opt['train']
self.DWT = common.DWT()
self.IWT = common.IWT()
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
# pretrained_dict = torch.load(opt['path']['pretrain_model_others'])
# netG_dict = self.netG.state_dict()
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in netG_dict}
# netG_dict.update(pretrained_dict)
# self.netG.load_state_dict(netG_dict)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
if self.is_train:
if not self.train_opt['only_G']:
self.netD = networks.define_D(opt).to(self.device)
# init_weights(self.netD)
if opt['dist']:
self.netD = DistributedDataParallel(
self.netD, device_ids=[torch.cuda.current_device()])
else:
self.netD = DataParallel(self.netD)
self.netG.train()
self.netD.train()
else:
self.netG.train()
else:
self.netG.train()
# define losses, optimizer and scheduler
if self.is_train:
# G pixel loss
if self.train_opt['pixel_weight'] > 0:
l_pix_type = self.train_opt['pixel_criterion']
if l_pix_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif l_pix_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
elif l_pix_type == 'cb':
self.cri_pix = CharbonnierLoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_pix_type))
self.l_pix_w = self.train_opt['pixel_weight']
else:
logger.info('Remove pixel loss.')
self.cri_pix = None
if self.train_opt['lpips_weight'] > 0:
l_lpips_type = self.train_opt['lpips_criterion']
if l_lpips_type == 'lpips':
self.cri_lpips = lpips.LPIPS(net='vgg').to(self.device)
if opt['dist']:
self.cri_lpips = DistributedDataParallel(
self.cri_lpips,
device_ids=[torch.cuda.current_device()])
else:
self.cri_lpips = DataParallel(self.cri_lpips)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(
l_lpips_type))
self.l_lpips_w = self.train_opt['lpips_weight']
else:
logger.info('Remove lpips loss.')
self.cri_lpips = None
# G feature loss
if self.train_opt['feature_weight'] > 0:
self.fea_trans = GramMatrix().to(self.device)
l_fea_type = self.train_opt['feature_criterion']
if l_fea_type == 'l1':
self.cri_fea = nn.L1Loss().to(self.device)
elif l_fea_type == 'l2':
self.cri_fea = nn.MSELoss().to(self.device)
elif l_fea_type == 'cb':
self.cri_fea = CharbonnierLoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_fea_type))
self.l_fea_w = self.train_opt['feature_weight']
else:
logger.info('Remove feature loss.')
self.cri_fea = None
if self.cri_fea: # load VGG perceptual loss
self.netF = networks.define_F(opt,
use_bn=False).to(self.device)
if opt['dist']:
self.netF = DistributedDataParallel(
self.netF, device_ids=[torch.cuda.current_device()])
else:
self.netF = DataParallel(self.netF)
# GD gan loss
self.cri_gan = GANLoss(self.train_opt['gan_type'], 1.0,
0.0).to(self.device)
self.l_gan_w = self.train_opt['gan_weight']
# D_update_ratio and D_init_iters
self.D_update_ratio = self.train_opt[
'D_update_ratio'] if self.train_opt['D_update_ratio'] else 1
self.D_init_iters = self.train_opt[
'D_init_iters'] if self.train_opt['D_init_iters'] else 0
# optimizers
# G
wd_G = self.train_opt['weight_decay_G'] if self.train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(
optim_params,
lr=self.train_opt['lr_G'],
weight_decay=wd_G,
betas=(self.train_opt['beta1_G'], self.train_opt['beta2_G']))
self.optimizers.append(self.optimizer_G)
if not self.train_opt['only_G']:
# D
wd_D = self.train_opt['weight_decay_D'] if self.train_opt[
'weight_decay_D'] else 0
self.optimizer_D = torch.optim.Adam(
self.netD.parameters(),
lr=self.train_opt['lr_D'],
weight_decay=wd_D,
betas=(self.train_opt['beta1_D'],
self.train_opt['beta2_D']))
self.optimizers.append(self.optimizer_D)
# schedulers
if self.train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
self.train_opt['lr_steps'],
restarts=self.train_opt['restarts'],
weights=self.train_opt['restart_weights'],
gamma=self.train_opt['lr_gamma'],
clear_state=self.train_opt['clear_state']))
elif self.train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
self.train_opt['T_period'],
eta_min=self.train_opt['eta_min'],
restarts=self.train_opt['restarts'],
weights=self.train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
if self.is_train:
if not self.train_opt['only_G']:
self.print_network() # print network
else:
self.print_network() # print network
try:
self.load() # load G and D if needed
print('Pretrained model loaded')
except Exception as e:
print('No pretrained model found')
def __init__(self, opt):
super(FIRNModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
test_opt = opt['test']
self.train_opt = train_opt
self.test_opt = test_opt
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
# print network
self.print_network()
self.load()
self.Quantization = Quantization()
if self.is_train:
self.netG.train()
# loss
self.Reconstruction_forw = ReconstructionLoss(
self.device, losstype=self.train_opt['pixel_criterion_forw'])
self.Reconstruction_back = ReconstructionLoss(
self.device, losstype=self.train_opt['pixel_criterion_back'])
# optimizers
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1'],
train_opt['beta2']))
self.optimizers.append(self.optimizer_G)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
def __init__(self, opt):
super(VideoBaseModel, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
# define network and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
# print network
self.print_network()
self.load()
if self.is_train:
self.netG.train()
#### loss
loss_type = train_opt['pixel_criterion']
if loss_type == 'l1':
self.cri_pix = nn.L1Loss(reduction='sum').to(self.device)
elif loss_type == 'l2':
self.cri_pix = nn.MSELoss(reduction='sum').to(self.device)
elif loss_type == 'cb':
self.cri_pix = CharbonnierLoss().to(self.device)
else:
raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type))
self.l_pix_w = train_opt['pixel_weight']
self.grad_w = train_opt['grad_weight']
#### optimizers
wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
if train_opt['ft_tsa_only']:
normal_params = []
tsa_fusion_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
if 'tsa_fusion' in k:
tsa_fusion_params.append(v)
else:
normal_params.append(v)
else:
if self.rank <= 0:
logger.warning('Params [{:s}] will not optimize.'.format(k))
optim_params = [
{ # add normal params first
'params': normal_params,
'lr': train_opt['lr_G']
},
{
'params': tsa_fusion_params,
'lr': train_opt['lr_G']
},
]
else:
optim_params = []
for k, v in self.netG.named_parameters():
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning('Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1'], train_opt['beta2']))
self.optimizers.append(self.optimizer_G)
#### schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
else:
raise NotImplementedError()
self.log_dict = OrderedDict()
def __init__(self, opt, edge_enhance=True):
super(SRGANModel, self).__init__(opt)
self.edge_enhance = edge_enhance
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
self.netG = networks.define_G(opt).to(self.device)
if opt['dist']:
self.netG = DistributedDataParallel(
self.netG, device_ids=[torch.cuda.current_device()])
else:
self.netG = DataParallel(self.netG)
if self.is_train:
self.netD = networks.define_D(opt).to(self.device)
if opt['dist']:
self.netD = DistributedDataParallel(
self.netD, device_ids=[torch.cuda.current_device()])
else:
self.netD = DataParallel(self.netD)
self.netG.train()
self.netD.train()
if self.is_train:
if train_opt['pixel_weight'] > 0:
l_pix_type = train_opt['pixel_criterion']
if l_pix_type == 'l1':
self.cri_pix = nn.L1Loss().to(self.device)
elif l_pix_type == 'l2':
self.cri_pix = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_pix_type))
self.l_pix_w = train_opt['pixel_weight']
else:
logger.info('Remove pixel loss.')
self.cri_pix = None
if train_opt['feature_weight'] > 0:
l_fea_type = train_opt['feature_criterion']
if l_fea_type == 'l1':
self.cri_fea = nn.L1Loss().to(self.device)
elif l_fea_type == 'l2':
self.cri_fea = nn.MSELoss().to(self.device)
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(l_fea_type))
self.l_fea_w = train_opt['feature_weight']
else:
logger.info('Remove feature loss.')
self.cri_fea = None
if self.cri_fea: # load VGG perceptual loss
self.netF = networks.define_F(opt,
use_bn=False).to(self.device)
if opt['dist']:
self.netF = DistributedDataParallel(
self.netF, device_ids=[torch.cuda.current_device()])
else:
self.netF = DataParallel(self.netF)
self.cri_gan = GANLoss(train_opt['gan_type'], 1.0,
0.0).to(self.device)
self.l_gan_w = train_opt['gan_weight']
self.D_update_ratio = train_opt['D_update_ratio'] if train_opt[
'D_update_ratio'] else 1
self.D_init_iters = train_opt['D_init_iters'] if train_opt[
'D_init_iters'] else 0
self.WGAN_QC_regul = QC_GradientPenaltyLoss()
if self.edge_enhance:
self.l_edge_w = train_opt['edge_weight']
if train_opt['edge_type'] == 'sobel':
self.cril_edge = sobel
elif train_opt['edge_type'] == 'canny':
self.cril_edge = canny
elif train_opt['edge_type'] == 'hednet':
self.netEdge = HedNet().cuda()
for p in self.netEdge.parameters():
p.requires_grad = False
self.cril_edge = self.netEdge
else:
raise NotImplementedError(
'Loss type [{:s}] not recognized.'.format(
train_opt['edge_type']))
else:
logger.info('Remove edge loss.')
self.cril_edge = None
wd_G = train_opt['weight_decay_G'] if train_opt[
'weight_decay_G'] else 0
optim_params = []
for k, v in self.netG.named_parameters(
): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
if self.rank <= 0:
logger.warning(
'Params [{:s}] will not optimize.'.format(k))
self.optimizer_G = torch.optim.Adam(optim_params,
lr=train_opt['lr_G'],
weight_decay=wd_G,
betas=(train_opt['beta1_G'],
train_opt['beta2_G']))
self.optimizers.append(self.optimizer_G)
wd_D = train_opt['weight_decay_D'] if train_opt[
'weight_decay_D'] else 0
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=train_opt['lr_D'],
weight_decay=wd_D,
betas=(train_opt['beta1_D'],
train_opt['beta2_D']))
self.optimizers.append(self.optimizer_D)
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.MultiStepLR_Restart(
optimizer,
train_opt['lr_steps'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights'],
gamma=train_opt['lr_gamma'],
clear_state=train_opt['clear_state']))
elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
for optimizer in self.optimizers:
self.schedulers.append(
lr_scheduler.CosineAnnealingLR_Restart(
optimizer,
train_opt['T_period'],
eta_min=train_opt['eta_min'],
restarts=train_opt['restarts'],
weights=train_opt['restart_weights']))
else:
raise NotImplementedError(
'MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
self.load()