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.netD = networks.define_D(opt).to(self.device)
self.netD_grad = networks.define_D_grad(opt).to(self.device) # D_grad
if opt['dist']:
self.netD = DistributedDataParallel(self.netD,
device_ids=[torch.cuda.current_device()],find_unused_parameters=True)
self.netD_grad = DistributedDataParallel(self.netD_grad,
device_ids=[torch.cuda.current_device()],find_unused_parameters=True)
else:
self.netD = DataParallel(self.netD)
self.netD_grad = DataParallel(self.netD_grad)
self.netG.train()
self.netD.train()
self.netD_grad.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
# 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)
# 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
# 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
if train_opt['gradient_pixel_weight'] > 0:
self.cri_pix_grad = nn.MSELoss().to(self.device)
self.l_pix_grad_w = train_opt['gradient_pixel_weight']
else:
self.cri_pix_grad = None
# gradient_gan_loss
if train_opt['gradient_gan_weight'] > 0:
self.cri_grad_gan = GANLoss(train_opt['gan_type'], 1.0, 0.0).to(self.device)
self.l_gan_grad_w = train_opt['gradient_gan_weight']
else:
self.cri_grad_gan = None
# 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)
# D_grad
wd_D_grad = train_opt['weight_decay_D'] if train_opt['weight_decay_D'] else 0
self.optimizer_D_grad = torch.optim.Adam(self.netD_grad.parameters(), lr=train_opt['lr_D'], \
weight_decay=wd_D, betas=(train_opt['beta1_D'], 0.999))
self.optimizers.append(self.optimizer_D_grad)
# 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(SPSRModel, self).__init__(opt)
train_opt = opt['train']
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device) # G
if self.is_train:
self.netD = networks.define_D(opt).to(self.device) # D
self.netD_grad = networks.define_D_grad(opt).to(self.device) # D_grad
self.netG.train()
self.netD.train()
self.netD_grad.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
# 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)
# 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 are for WGAN
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
# 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
if train_opt['gan_type'] == 'wgan-gp':
self.random_pt = torch.Tensor(1, 1, 1, 1).to(self.device)
# gradient penalty loss
self.cri_gp = GradientPenaltyLoss(device=self.device).to(self.device)
self.l_gp_w = train_opt['gp_weigth']
# gradient_pixel_loss
if train_opt['gradient_pixel_weight'] > 0:
self.cri_pix_grad = nn.MSELoss().to(self.device)
self.l_pix_grad_w = train_opt['gradient_pixel_weight']
else:
self.cri_pix_grad = None
# gradient_gan_loss
if train_opt['gradient_gan_weight'] > 0:
self.cri_grad_gan = GANLoss(train_opt['gan_type'], 1.0, 0.0).to(self.device)
self.l_gan_grad_w = train_opt['gradient_gan_weight']
else:
self.cri_grad_gan = None
# G_grad pixel loss
if train_opt['pixel_branch_weight'] > 0:
l_pix_type = train_opt['pixel_branch_criterion']
if l_pix_type == 'l1':
self.cri_pix_branch = nn.L1Loss().to(self.device)
elif l_pix_type == 'l2':
self.cri_pix_branch = nn.MSELoss().to(self.device)
else:
raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_pix_type))
self.l_pix_branch_w = train_opt['pixel_branch_weight']
else:
logger.info('Remove G_grad pixel loss.')
self.cri_pix_branch = None
# 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(): # optimize part of the model
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=train_opt['lr_G'], \
weight_decay=wd_G, betas=(train_opt['beta1_G'], 0.999))
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'], 0.999))
self.optimizers.append(self.optimizer_D)
# D_grad
wd_D_grad = train_opt['weight_decay_D'] if train_opt['weight_decay_D'] else 0
self.optimizer_D_grad = torch.optim.Adam(self.netD_grad.parameters(), lr=train_opt['lr_D'], \
weight_decay=wd_D, betas=(train_opt['beta1_D'], 0.999))
self.optimizers.append(self.optimizer_D_grad)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \
train_opt['lr_steps'], train_opt['lr_gamma']))
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()