def configure_loggers(opt=None):
tofile = opt.get('logger', {}).get('save_logfile', True)
if opt['is_train']:
# config loggers. Before it, the log will not work
util.get_root_logger(None,
opt['path']['log'],
'train',
level=logging.INFO,
screen=True,
tofile=tofile)
util.get_root_logger('val',
opt['path']['log'],
'val',
level=logging.INFO,
tofile=tofile)
else:
util.get_root_logger(None,
opt['path']['log'],
'test',
level=logging.INFO,
screen=True,
tofile=tofile)
logger = util.get_root_logger() # 'base'
logger.info(options.dict2str(opt))
# initialize tensorboard logger
tb_logger = None
if opt.get('use_tb_logger', False) and 'debug' not in opt['name']:
version = float(torch.__version__[0:3])
log_dir = os.path.join(opt['path']['root'], 'tb_logger', opt['name'])
# log_dir = os.path.join(opt['path']['experiments_root'], opt['name'], 'tb')
# logdir_valid = os.path.join(opt['path']['root'], 'tb_logger', opt['name'] + 'valid')
# logdir_valid = os.path.join(opt['path']['experiments_root'], opt['name'], 'tb_valid')
if version >= 1.1: # PyTorch 1.1
# official PyTorch tensorboard
try:
from torch.utils.tensorboard import SummaryWriter
except:
from tensorboardX import SummaryWriter
else:
logger.info(
'You are using PyTorch {}. Using [tensorboardX].'.format(
version))
from tensorboardX import SummaryWriter
try:
# for versions PyTorch > 1.1 and tensorboardX < 1.6
tb_logger = SummaryWriter(log_dir=log_dir)
# tb_logger_valid = SummaryWriter(log_dir=logdir_valid)
except:
# for version tensorboardX >= 1.7
tb_logger = SummaryWriter(logdir=log_dir)
# tb_logger_valid = SummaryWriter(logdir=logdir_valid)
return {"tb_logger": tb_logger}
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, required=True, help='Path to options file.')
opt = options.parse(parser.parse_args().opt, is_train=False)
util.mkdirs((path for key, path in opt['path'].items() if not key == 'pretrain_model_G'))
logger = util.get_root_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True)
logger = logging.getLogger('base')
logger.info(options.dict2str(opt))
scale = opt.get('scale', 4)
# Create test dataset and dataloader
test_loaders = []
znorm = False # TMP
# znorm_list = []
'''
video_list = os.listdir(cfg.testset_dir)
for idx_video in range(len(video_list)):
video_name = video_list[idx_video]
# dataloader
test_set = TestsetLoader(cfg, video_name)
test_loader = DataLoader(test_set, num_workers=1, batch_size=1, shuffle=False)
'''
for phase, dataset_opt in sorted(opt['datasets'].items()):
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set, dataset_opt)
logger.info('Number of test images in [{:s}]: {:d}'.format(dataset_opt['name'], len(test_set)))
test_loaders.append(test_loader)
# Temporary, will turn znorm on for all the datasets. Will need to introduce a variable for each dataset and differentiate each one later in the loop.
# if dataset_opt.get['znorm'] and znorm == False:
# znorm = True
znorm = dataset_opt.get('znorm', False)
# znorm_list.apped(znorm)
# Create model
model = create_model(opt)
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
logger.info('\nTesting [{:s}]...'.format(test_set_name))
test_start_time = time.time()
dataset_dir = os.path.join(opt['path']['results_root'], test_set_name)
util.mkdir(dataset_dir)
test_results = OrderedDict()
test_results['psnr'] = []
test_results['ssim'] = []
test_results['psnr_y'] = []
test_results['ssim_y'] = []
for data in test_loader:
need_HR = False if test_loader.dataset.opt['dataroot_HR'] is None else True
img_path = data['LR_path'][0]
img_name = os.path.splitext(os.path.basename(img_path))[0]
# tmp_vis(data['LR'][:,1,:,:,:], True)
if opt.get('chop_forward', None):
# data
if len(data['LR'].size()) == 4:
b, n_frames, h_lr, w_lr = data['LR'].size()
LR_y_cube = data['LR'].view(b, -1, 1, h_lr, w_lr) # b, t, c, h, w
elif len(data['LR'].size()) == 5: # for networks that work with 3 channel images
_, n_frames, _, _, _ = data['LR'].size()
LR_y_cube = data['LR'] # b, t, c, h, w
# print(LR_y_cube.shape)
# print(data['LR_bicubic'].shape)
# crop borders to ensure each patch can be divisible by 2
# TODO: this is modcrop, not sure if really needed, check (the dataloader already does modcrop)
_, _, _, h, w = LR_y_cube.size()
h = int(h // 16) * 16
w = int(w // 16) * 16
LR_y_cube = LR_y_cube[:, :, :, :h, :w]
if isinstance(data['LR_bicubic'], torch.Tensor):
# SR_cb = data['LR_bicubic'][:, 1, :, :][:, :, :h * scale, :w * scale]
SR_cb = data['LR_bicubic'][:, 1, :h * scale, :w * scale]
# SR_cr = data['LR_bicubic'][:, 2, :, :][:, :, :h * scale, :w * scale]
SR_cr = data['LR_bicubic'][:, 2, :h * scale, :w * scale]
SR_y = chop_forward(LR_y_cube, model, scale, need_HR=need_HR).squeeze(0)
# SR_y = np.array(SR_y.data.cpu())
if test_loader.dataset.opt.get('srcolors', None):
print(SR_y.shape, SR_cb.shape, SR_cr.shape)
sr_img = ycbcr_to_rgb(torch.stack((SR_y, SR_cb, SR_cr), -3))
else:
sr_img = SR_y
else:
# data
model.feed_data(data, need_HR=need_HR)
# SR_y = net(LR_y_cube).squeeze(0)
model.test() # test
visuals = model.get_current_visuals(need_HR=need_HR)
# ds = torch.nn.AvgPool2d(2, stride=2, count_include_pad=False)
# tmp_vis(ds(visuals['SR']), True)
# tmp_vis(visuals['SR'], True)
if test_loader.dataset.opt.get('y_only', None) and test_loader.dataset.opt.get('srcolors', None):
SR_cb = data['LR_bicubic'][:, 1, :, :]
SR_cr = data['LR_bicubic'][:, 2, :, :]
# tmp_vis(ds(SR_cb), True)
# tmp_vis(ds(SR_cr), True)
sr_img = ycbcr_to_rgb(torch.stack((visuals['SR'], SR_cb, SR_cr), -3))
else:
sr_img = visuals['SR']
# if znorm the image range is [-1,1], Default: Image range is [0,1] # testing, each "dataset" can have a different name (not train, val or other)
sr_img = tensor2np(sr_img, denormalize=znorm) # uint8
# save images
suffix = opt['suffix']
if suffix:
save_img_path = os.path.join(dataset_dir, img_name + suffix + '.png')
else:
save_img_path = os.path.join(dataset_dir, img_name + '.png')
util.save_img(sr_img, save_img_path)
# TODO: update to use metrics functions
# calculate PSNR and SSIM
if need_HR:
# if znorm the image range is [-1,1], Default: Image range is [0,1] # testing, each "dataset" can have a different name (not train, val or other)
gt_img = tensor2img(visuals['HR'], denormalize=znorm) # uint8
gt_img = gt_img / 255.
sr_img = sr_img / 255.
crop_border = test_loader.dataset.opt['scale']
cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :]
cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :]
psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255)
ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255)
test_results['psnr'].append(psnr)
test_results['ssim'].append(ssim)
if gt_img.shape[2] == 3: # RGB image
sr_img_y = bgr2ycbcr(sr_img, only_y=True)
gt_img_y = bgr2ycbcr(gt_img, only_y=True)
cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border]
cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border]
psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255)
ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255)
test_results['psnr_y'].append(psnr_y)
test_results['ssim_y'].append(ssim_y)
logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.' \
.format(img_name, psnr, ssim, psnr_y, ssim_y))
else:
logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}.'.format(img_name, psnr, ssim))
else:
logger.info(img_name)
# TODO: update to use metrics functions
if need_HR: # metrics
# Average PSNR/SSIM results
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
logger.info('----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n' \
.format(test_set_name, ave_psnr, ave_ssim))
if test_results['psnr_y'] and test_results['ssim_y']:
ave_psnr_y = sum(test_results['psnr_y']) / len(test_results['psnr_y'])
ave_ssim_y = sum(test_results['ssim_y']) / len(test_results['ssim_y'])
logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n' \
.format(ave_psnr_y, ave_ssim_y))
def main():
#### options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, help='Path to option YMAL file.')
parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
help='job launcher')
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
opt = option.parse(args.opt, is_train=True)
#### distributed training settings
if args.launcher == 'none': # disabled distributed training
opt['dist'] = False
rank = -1
print('Disabled distributed training.')
else:
opt['dist'] = True
init_dist()
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
#### loading resume state if exists
if opt['path'].get('resume_state', None):
# distributed resuming: all load into default GPU
device_id = torch.cuda.current_device()
resume_state = torch.load(opt['path']['resume_state'],
map_location=lambda storage, loc: storage.cuda(device_id))
option.check_resume(opt, resume_state['iter']) # check resume options
else:
resume_state = None
#### mkdir and loggers
if rank <= 0: # normal training (rank -1) OR distributed training (rank 0)
if resume_state is None:
print(opt['path'])
util.mkdir_and_rename(
opt['path']['experiments_root']) # rename experiment folder if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root'
and 'pretrain_model' not in key and 'resume' not in key and path is not None))
# config loggers. Before it, the log will not work
util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO,
screen=True, tofile=True)
util.setup_logger('val', opt['path']['log'], 'val_' + opt['name'], level=logging.INFO,
screen=True, tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
version = float(torch.__version__[0:3])
if version >= 1.1: # PyTorch 1.1
from torch.utils.tensorboard import SummaryWriter
else:
logger.info(
'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version))
from tensorboardX import SummaryWriter
trial = 0
while os.path.isdir('../Loggers/' + opt['name'] + '/' + str(trial)):
trial += 1
tb_logger = SummaryWriter(log_dir='../Loggers/' + opt['name'] + '/' + str(trial))
else:
util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True)
logger = logging.getLogger('base')
# convert to NoneDict, which returns None for missing keys
opt = option.dict_to_nonedict(opt)
# -------------------------------------------- ADDED --------------------------------------------
l1_loss = torch.nn.L1Loss()
mse_loss = torch.nn.MSELoss()
calc_lpips = PerceptualLossLPIPS()
if torch.cuda.is_available():
l1_loss = l1_loss.cuda()
mse_loss = mse_loss.cuda()
# -----------------------------------------------------------------------------------------------
#### random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
if rank <= 0:
logger.info('Random seed: {}'.format(seed))
util.set_random_seed(seed)
torch.backends.cudnn.benckmark = True
# torch.backends.cudnn.deterministic = True
#### create train and val dataloader
dataset_ratio = 200 # enlarge the size of each epoch
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
train_set = create_dataset(dataset_opt)
train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size']))
total_iters = int(opt['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
if opt['dist']:
train_sampler = DistIterSampler(train_set, world_size, rank, dataset_ratio)
total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio)))
else:
train_sampler = None
train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler)
if rank <= 0:
logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
total_epochs, total_iters))
elif phase == 'val':
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt, opt, None)
if rank <= 0:
logger.info('Number of val images in [{:s}]: {:d}'.format(
dataset_opt['name'], len(val_set)))
else:
raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
assert train_loader is not None
#### create model
model = Model(opt)
#### resume training
if resume_state:
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
start_epoch = resume_state['epoch']
current_step = resume_state['iter']
model.resume_training(resume_state) # handle optimizers and schedulers
else:
current_step = 0
start_epoch = 0
#### training
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
for epoch in range(start_epoch, total_epochs + 1):
if opt['dist']:
train_sampler.set_epoch(epoch)
train_bar = tqdm(train_loader, desc='[%d/%d]' % (epoch, total_epochs))
for bus, train_data in enumerate(train_bar):
# validation
if epoch % opt['train']['val_freq'] == 0 and bus == 0 and rank <= 0:
avg_ssim = avg_psnr = avg_lpips = val_pix_err_f = val_pix_err_nf = val_mean_color_err = 0.0
print("into validation!")
idx = 0
val_bar = tqdm(val_loader, desc='[%d/%d]' % (epoch, total_epochs))
for val_data in val_bar:
idx += 1
img_name = os.path.splitext(os.path.basename(val_data['LQ_path'][0]))[0]
img_dir = os.path.join(opt['path']['val_images'], img_name)
util.mkdir(img_dir)
model.feed_data(val_data)
model.test()
visuals = model.get_current_visuals()
sr_img = util.tensor2img(visuals['SR']) # uint8
gt_img = util.tensor2img(visuals['GT']) # uint8
lq_img = util.tensor2img(visuals['LQ']) # uint8
#nr_img = util.tensor2img(visuals['NR']) # uint8
#nf_img = util.tensor2img(visuals['NF']) # uint8
#nh_img = util.tensor2img(visuals['NH']) # uint8
#print("Great! images got into here.")
# Save SR images for reference
save_sr_img_path = os.path.join(img_dir,
'{:s}_{:d}_sr.png'.format(img_name, current_step))
save_nr_img_path = os.path.join(img_dir,
'{:s}_{:d}_lq.png'.format(img_name, current_step))
#save_nf_img_path = os.path.join(img_dir,
# 'bs_{:s}_{:d}_nr.png'.format(img_name, current_step))
#save_nh_img_path = os.path.join(img_dir,
# 'bs_{:s}_{:d}_nh.png'.format(img_name, current_step))
util.save_img(sr_img, save_sr_img_path)
util.save_img(lq_img, save_nr_img_path)
#util.save_img(nf_img, save_nf_img_path)
#util.save_img(nh_img, save_nh_img_path)
#print("Saved")
# calculate PSNR
gt_img = gt_img / 255.
sr_img = sr_img / 255.
#nf_img = nf_img / 255.
lq_img = lq_img / 255.
#cropped_lq_img = lq_img[crop_size:-crop_size, crop_size:-crop_size, :]
#cropped_nr_img = nr_img[crop_size:-crop_size, crop_size:-crop_size, :]
avg_psnr += util.calculate_psnr(sr_img * 255, gt_img * 255)
avg_ssim += util.calculate_ssim(sr_img * 255, gt_img * 255)
avg_lpips += calc_lpips(visuals['SR'], visuals['GT'])
#avg_psnr_n += util.calculate_psnr(cropped_lq_img * 255, cropped_nr_img * 255)
# ----------------------------------------- ADDED -----------------------------------------
val_pix_err_nf += l1_loss(visuals['SR'], visuals['GT'])
val_mean_color_err += mse_loss(visuals['SR'].mean(2).mean(1), visuals['GT'].mean(2).mean(1))
# -----------------------------------------------------------------------------------------
avg_psnr = avg_psnr / idx
avg_ssim = avg_ssim / idx
avg_lpips = avg_lpips / idx
val_pix_err_f /= idx
val_pix_err_nf /= idx
val_mean_color_err /= idx
# log
logger.info('# Validation # PSNR: {:.4e},'.format(avg_psnr))
logger.info('# Validation # SSIM: {:.4e},'.format(avg_ssim))
logger.info('# Validation # LPIPS: {:.4e},'.format(avg_lpips))
logger_val = logging.getLogger('val') # validation logger
logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e} ssim: {:.4e} lpips: {:.4e}'.format(
epoch, current_step, avg_psnr, avg_ssim, avg_lpips))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('val_psnr', avg_psnr, current_step)
tb_logger.add_scalar('val_ssim', avg_ssim, current_step)
tb_logger.add_scalar('val_lpips', avg_lpips, current_step)
tb_logger.add_scalar('val_pix_err_nf', val_pix_err_nf, current_step)
tb_logger.add_scalar('val_mean_color_err', val_mean_color_err, current_step)
current_step += 1
if current_step > total_iters:
break
#### update learning rate
model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter'])
#### training
model.feed_data(train_data)
model.optimize_parameters(current_step)
model.clear_data()
#### tb_logger
if current_step % opt['logger']['tb_freq'] == 0:
logs = model.get_current_log()
if opt['use_tb_logger'] and 'debug' not in opt['name']:
for k, v in logs.items():
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
#### logger
if epoch % opt['logger']['print_freq'] == 0 and epoch != 0 and bus == 0:
logs = model.get_current_log()
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step, model.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
if rank <= 0:
logger.info(message)
#### save models and training states
if epoch % opt['logger']['save_checkpoint_freq'] == 0 and epoch != 0 and bus == 0:
if rank <= 0:
logger.info('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
if rank <= 0:
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training.')