def eval_path(args): if not os.path.exists(args.outputs_dir): os.makedirs(args.outputs_dir) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) cudnn.benchmark = True device_ids = list(range(args.gpus)) model = EDSR(args) model = nn.DataParallel(model, device_ids=device_ids) model = model.cuda() 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) 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) else: print("=> no checkpoint found at '{}'".format(args.resume)) model.eval() file_names = sorted(os.listdir(args.test_lr)) lr_list = [] for one in file_names: dst_dir = os.path.join(args.outputs_dir, one) if os.path.exists(dst_dir) and len(os.listdir(dst_dir)) == 100: continue lr_tmp = sorted(glob(os.path.join(args.test_lr, one, '*.png'))) lr_list.extend(lr_tmp) data_set = EvalDataset(lr_list) eval_loader = DataLoader(data_set, batch_size=args.batch_size, num_workers=args.workers) with tqdm(total=(len(data_set) - len(data_set) % args.batch_size)) as t: for data in eval_loader: inputs, names = data inputs = inputs.cuda() with torch.no_grad(): outputs = model(inputs).data.float().cpu().clamp_(0, 255).numpy() for img, file in zip(outputs, names): img = np.transpose(img[[2, 1, 0], :, :], (1, 2, 0)) img = img.round() arr = file.split('/') dst_dir = os.path.join(args.outputs_dir, arr[-2]) if not os.path.exists(dst_dir): os.makedirs(dst_dir) dst_name = os.path.join(dst_dir, arr[-1]) cv2.imwrite(dst_name, img) t.update(len(names))
def main(args): print("===> Loading datasets") file_name = sorted(os.listdir(args.data_lr)) lr_list = [] hr_list = [] for one in file_name: lr_tmp = sorted(glob(os.path.join(args.data_lr, one, '*.png'))) lr_list.extend(lr_tmp) hr_tmp = sorted(glob(os.path.join(args.data_hr, one, '*.png'))) if len(hr_tmp) != 100: print(one) hr_list.extend(hr_tmp) # lr_list = glob(os.path.join(args.data_lr, '*')) # hr_list = glob(os.path.join(args.data_hr, '*')) # lr_list = lr_list[0:max_index] # hr_list = hr_list[0:max_index] data_set = DatasetLoader(lr_list, hr_list, args.patch_size, args.scale) data_len = len(data_set) train_loader = DataLoader(data_set, batch_size=args.batch_size, num_workers=args.workers, shuffle=True, pin_memory=True, drop_last=True) print("===> Building model") np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) cudnn.benchmark = True device_ids = list(range(args.gpus)) model = EDSR(args) criterion = nn.L1Loss(reduction='sum') print("===> Setting GPU") model = nn.DataParallel(model, device_ids=device_ids) model = model.cuda() criterion = criterion.cuda() 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) if args.start_epoch != 0: #如果设置了 start_epoch 则不用checkpoint中的epoch参数 start_epoch = args.start_epoch print("===> Setting Optimizer") optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.999), eps=1e-08) # record = [] print("===> Training") for epoch in range(start_epoch, args.epochs): adjust_lr(optimizer, epoch) losses, psnrs = one_epoch_train_logger(model, optimizer, criterion, data_len, train_loader, epoch, args.epochs, args.batch_size, optimizer.param_groups[0]["lr"]) # lr = optimizer.param_groups[0]["lr"] # the_lr = 1e-2 # lr_len = 2 # while lr + (1e-9) < the_lr: # the_lr *= 0.1 # lr_len += 1 # record.append([losses.avg,psnrs.avg,lr_len]) # save model # if epoch+1 != args.epochs: # continue model_out_path = os.path.join( args.checkpoint, "model_epoch_%04d_edsr_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)