def __init__(self, hparams: Namespace) -> None: super(SR3Experiment, self).__init__() self.hparams = hparams self.val_folder = Path('output') self.val_folder.mkdir(exist_ok=True, parents=True) self.create_model() self.scale = self.hparams.scale # loss loss_type = self.hparams.train.pixel_criterion self.loss_type = loss_type if loss_type == 'l1': self.cri_pix = nn.L1Loss(reduction='none') elif loss_type == 'l2': self.cri_pix = nn.MSELoss(reduction='none') elif loss_type == 'cb': self.cri_pix = CharbonnierLoss() elif loss_type == 'l1l2': class L1L2Loss(nn.Module): def __init__(self, reduction='mean'): super().__init__() self.l1 = nn.L1Loss(reduction=reduction) self.l2 = nn.MSELoss(reduction=reduction) def forward(self, x, y): return self.l1(x, y) + self.l2(x, y) self.cri_pix = L1L2Loss(reduction='none') else: raise NotImplementedError( 'Loss type [{:s}] is not recognized.'.format(loss_type)) self.l_pix_w = self.hparams.train.pixel_weight
def get_criterion(self, mode, opt): if mode == 'pix': loss_type = opt['pixel_criterion'] if loss_type == 'l1': criterion = nn.L1Loss(reduction=opt['reduction']).to(self.device) elif loss_type == 'l2': criterion = nn.MSELoss(reduction=opt['reduction']).to(self.device) elif loss_type == 'cb': criterion = CharbonnierLoss(reduction=opt['reduction']).to(self.device) else: raise NotImplementedError('Loss type [{:s}] is not recognized for pixel'. format(loss_type)) weight = opt['pixel_weight'] elif mode == 'gan': criterion = GANLoss(opt['gan_type'], 1.0, 0.0).to(self.device) weight = opt['gan_weight'] else: raise TypeError('Unknown type: {} for criterion'.format(mode)) return criterion, weight
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): 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(SRModel, 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().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'] # Gradient Loss if train_opt['gradient_weight'] > 0: self.cri_gradient = GradientLoss(train_opt['gradient_type'], train_opt['gradient_grid_horizontal'], train_opt['gradient_grid_vertical'], train_opt['gradient_criterion'], self.device) self.l_gradient_w = train_opt['gradient_weight'] else: logger.info('Remove gradient loss') self.cri_gradient = None # 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.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 main(args): print("===> Loading datasets") data_set = DatasetLoader(args.data_lr, args.data_hr, size_w=args.size_w, size_h=args.size_h, scale=args.scale, n_frames=args.n_frames, interval_list=args.interval_list, border_mode=args.border_mode, random_reverse=args.random_reverse) train_loader = DataLoader(data_set, batch_size=args.batch_size, num_workers=args.workers, shuffle=True, pin_memory=False, drop_last=True) #### random seed random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) cudnn.benchmark = True #cudnn.deterministic = True print("===> Building model") #### create model model = EDVR_arch.EDVR(nf=args.nf, nframes=args.n_frames, groups=args.groups, front_RBs=args.front_RBs, back_RBs=args.back_RBs, center=args.center, predeblur=args.predeblur, HR_in=args.HR_in, w_TSA=args.w_TSA) criterion = CharbonnierLoss() print("===> Setting GPU") gups = args.gpus if args.gpus != 0 else torch.cuda.device_count() device_ids = list(range(gups)) model = DataParallel(model, device_ids=device_ids) model = model.cuda() criterion = criterion.cuda() # print(model) start_epoch = args.start_epoch # optionally resume from a checkpoint if args.resume: if os.path.isdir(args.resume): # 获取目录中最后一个 pth_list = sorted(glob(os.path.join(args.resume, '*.pth'))) if len(pth_list) > 0: args.resume = pth_list[-1] if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) start_epoch = checkpoint['epoch'] + 1 state_dict = checkpoint['state_dict'] new_state_dict = OrderedDict() for k, v in state_dict.items(): namekey = 'module.' + k # remove `module.` new_state_dict[namekey] = v model.load_state_dict(new_state_dict) # 如果文件中有lr,则不用启动参数 args.lr = checkpoint.get('lr', args.lr) # 如果设置了 start_epoch 则不用checkpoint中的epoch参数 start_epoch = args.start_epoch if args.start_epoch != 0 else start_epoch #如果use_current_lr大于0 测代替作为lr args.lr = args.use_current_lr if args.use_current_lr > 0 else args.lr optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay, betas=(args.beta1, args.beta2), eps=1e-8) #### training print("===> Training") for epoch in range(start_epoch, args.epochs): adjust_lr(optimizer, epoch) if args.use_tqdm == 1: losses, psnrs = one_epoch_train_tqdm( model, optimizer, criterion, len(data_set), train_loader, epoch, args.epochs, args.batch_size, optimizer.param_groups[0]["lr"]) else: losses, psnrs = one_epoch_train_logger( model, optimizer, criterion, len(data_set), train_loader, epoch, args.epochs, args.batch_size, optimizer.param_groups[0]["lr"]) # save model # if epoch %9 != 0: # continue model_out_path = os.path.join( args.checkpoint, "model_epoch_%04d_edvr_loss_%.3f_psnr_%.3f.pth" % (epoch, losses.avg, psnrs.avg)) if not os.path.exists(args.checkpoint): os.makedirs(args.checkpoint) torch.save( { 'state_dict': model.module.state_dict(), "epoch": epoch, 'lr': optimizer.param_groups[0]["lr"] }, model_out_path)
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()