def main():
"""The main function."""
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option JSON file.')
options = option.parse(
parser.parse_args().opt) # load settings and initialize settings
util.mkdir_and_rename(
options['path']['results_root']) # rename old experiments if exists
util.mkdirs((path for key, path in options['path'].items()
if not key == 'experiments_root' and not key == 'saved_model'
and not key == 'fc_root' and not key == 'db_root'))
option.save(options)
options = option.dict_to_nonedict(
options) # Convert to NoneDict, which return None for missing key.
# Redirect all writes to the "txt" file
sys.stdout = PrintLogger(options['path']['log'])
# # logger = Logger(opt)
if options['manual_seed'] is not None:
random.seed(options['manual_seed'])
index = None
if options['dataset'] == 'pipal':
index = list(range(0, 200))
else:
raise NotImplementedError
lr_backup = options['base_lr']
lr_steps_backup = options['lr_steps']
epochs_backup = options['epochs']
srcc_all = np.zeros((1, options['epoch']), dtype=np.float)
plcc_all = np.zeros((1, options['epoch']), dtype=np.float)
print('Total epochs:' + str(options['epoch']))
for i in range(0, options['epoch']):
# randomly split train-test set
random.shuffle(index)
train_index = index[0:round(0.8 * len(index))]
test_index = index[round(0.8 * len(index)):len(index)]
options['train_index'] = train_index
options['test_index'] = test_index
# train the fully connected layer only
print('[No.%d/%d] Training the FC layer...' % (i, options['epoch']))
options['fc'] = True
options['base_lr'] = lr_backup
options['lr_steps'] = lr_steps_backup
options['epochs'] = epochs_backup
manager = IQAManager(options)
best_srcc1, best_plcc1 = manager.train()
# fine-tune all model
print('[No.%d/%d] Fine-tune all model...' % (i, options['epoch']))
options['fc'] = False
options['base_lr'] = options['base_lr_step2']
options['lr_steps'] = options['lr_steps_step2']
options['epochs'] = options['epochs_step2']
manager = IQAManager(options)
best_srcc2, best_plcc2 = manager.train()
srcc_all[0][i] = np.max([best_srcc1, best_srcc2])
plcc_all[0][i] = np.max([best_plcc1, best_plcc2])
# srcc_all[0][i] = best_srcc1
# plcc_all[0][i] = best_plcc1
srcc_mean = np.mean(srcc_all)
srcc_median = np.median(srcc_all)
plcc_mean = np.mean(plcc_all)
plcc_median = np.median(plcc_all)
print(srcc_all)
print('average srcc:%4.4f' % srcc_mean)
print('median srcc:%4.4f' % srcc_median)
print(plcc_all)
print('average plcc:%4.4f' % plcc_mean)
print('median plcc:%4.4f' % plcc_median)
print('--------------Finish! [' + options['name'] + ']--------------')
util.mkdir_and_rename(options['path']['results_root'],
('done_{:.4f}'.format(srcc_median)))
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=True)
util.mkdir_and_rename(
opt['path']['experiments_root']) # rename old experiments if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and \
not key == 'pretrain_model_G' and not key == 'pretrain_model_D'))
option.save(opt)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# print to file and std_out simultaneously
sys.stdout = PrintLogger(opt['path']['log'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
print("Random Seed: ", seed)
random.seed(seed)
torch.manual_seed(seed)
# create train and val dataloader
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']))
print('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epoches = int(math.ceil(total_iters / train_size))
print('Total epoches needed: {:d} for iters {:,d}'.format(
total_epoches, total_iters))
train_loader = create_dataloader(train_set, dataset_opt)
elif phase == 'val':
val_dataset_opt = dataset_opt
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
print('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 = create_model(opt)
# create logger
logger = Logger(opt)
current_step = 0
start_time = time.time()
print('---------- Start training -------------')
for epoch in range(total_epoches):
for i, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
# training
model.feed_data(train_data)
model.optimize_parameters(current_step)
time_elapsed = time.time() - start_time
start_time = time.time()
# log
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
print_rlt['time'] = time_elapsed
for k, v in logs.items():
print_rlt[k] = v
print_rlt['lr'] = model.get_current_learning_rate()
logger.print_format_results('train', print_rlt)
# save models
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
print('Saving the model at the end of iter {:d}.'.format(
current_step))
model.save(current_step)
# validation
if current_step % opt['train']['val_freq'] == 0:
print('---------- validation -------------')
start_time = time.time()
avg_psnr = 0.0
idx = 0
for val_data in val_loader:
idx += 1
img_name = os.path.splitext(
os.path.basename(val_data['LR_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['HR']) # uint8
# Save SR images for reference
save_img_path = os.path.join(img_dir, '{:s}_{:d}.png'.format(\
img_name, current_step))
util.save_img(sr_img, save_img_path)
# calculate PSNR
if opt['crop_scale'] is not None:
crop_size = opt['crop_scale']
else:
crop_size = opt['scale']
if crop_size <= 0:
cropped_sr_img = sr_img.copy()
cropped_gt_img = gt_img.copy()
else:
if len(gt_img.shape) < 3:
cropped_sr_img = sr_img[crop_size:-crop_size,
crop_size:-crop_size]
cropped_gt_img = gt_img[crop_size:-crop_size,
crop_size:-crop_size]
else:
cropped_sr_img = sr_img[crop_size:-crop_size,
crop_size:-crop_size, :]
cropped_gt_img = gt_img[crop_size:-crop_size,
crop_size:-crop_size, :]
#avg_psnr += util.psnr(cropped_sr_img, cropped_gt_img)
cropped_sr_img_y = bgr2ycbcr(cropped_sr_img, only_y=True)
cropped_gt_img_y = bgr2ycbcr(cropped_gt_img, only_y=True)
avg_psnr += util.psnr(
cropped_sr_img_y,
cropped_gt_img_y) ##########only y channel
avg_psnr = avg_psnr / idx
time_elapsed = time.time() - start_time
# Save to log
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
print_rlt['time'] = time_elapsed
print_rlt['psnr'] = avg_psnr
logger.print_format_results('val', print_rlt)
print('-----------------------------------')
# update learning rate
model.update_learning_rate()
print('Saving the final model.')
model.save('latest')
print('End of training.')
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=True)
util.mkdir_and_rename(
opt['path']['experiments_root']) # rename old experiments if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and \
not key == 'pretrain_model_G' and not key == 'pretrain_model_D'))
option.save(opt)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# print to file and std_out simultaneously
sys.stdout = PrintLogger(opt['path']['log'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
print("Random Seed: ", seed)
random.seed(seed)
torch.manual_seed(seed)
# create train and val dataloader
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']))
print('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epoches = int(math.ceil(total_iters / train_size))
print('Total epoches needed: {:d} for iters {:,d}'.format(
total_epoches, total_iters))
train_loader = create_dataloader(train_set, dataset_opt)
batch_size_per_month = dataset_opt['batch_size']
batch_size_per_day = int(
opt['datasets']['train']['batch_size_per_day'])
num_month = int(opt['train']['num_month'])
num_day = int(opt['train']['num_day'])
elif phase == 'val':
val_dataset_opt = dataset_opt
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
print('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 = create_model(opt)
# create logger
logger = Logger(opt)
current_step = 0
start_time = time.time()
print('---------- Start training -------------')
validate(val_loader, opt, model, current_step, 0, logger)
for epoch in range(num_month):
for i, train_data in enumerate(train_loader):
cur_month_code = get_code_for_data_two(model, train_data, opt)
for j in range(num_day):
current_step += 1
# get the sliced data
cur_day_batch_start_idx = (
j * batch_size_per_day) % batch_size_per_month
cur_day_batch_end_idx = cur_day_batch_start_idx + batch_size_per_day
if cur_day_batch_end_idx > batch_size_per_month:
cur_day_batch_idx = np.hstack(
(np.arange(cur_day_batch_start_idx,
batch_size_per_month),
np.arange(cur_day_batch_end_idx -
batch_size_per_month)))
else:
cur_day_batch_idx = slice(cur_day_batch_start_idx,
cur_day_batch_end_idx)
cur_day_train_data = {
'LR': train_data['LR'][cur_day_batch_idx],
'HR': train_data['HR'][cur_day_batch_idx]
}
code = []
for gen_code in cur_month_code:
code.append(gen_code[cur_day_batch_idx])
# training
model.feed_data(cur_day_train_data, code=code)
model.optimize_parameters(current_step)
time_elapsed = time.time() - start_time
start_time = time.time()
# log
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
print_rlt['time'] = time_elapsed
for k, v in logs.items():
print_rlt[k] = v
print_rlt['lr'] = model.get_current_learning_rate()
logger.print_format_results('train', print_rlt)
# save models
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
print('Saving the model at the end of iter {:d}.'.format(
current_step))
model.save(current_step)
# validation
if current_step % opt['train']['val_freq'] == 0:
validate(val_loader, opt, model, current_step, epoch,
logger)
# update learning rate
model.update_learning_rate()
print('Saving the final model.')
model.save('latest')
print('End of training.')
parser.add_argument('-single_GPU',
action='store_true',
help='Utilize only one GPU')
if parser.parse_args().single_GPU:
util.Assign_GPU()
opt = option.parse(parser.parse_args().opt, is_train=False)
util.mkdirs((path for key, path in opt['path'].items()
if not key == 'pretrained_model_G'))
opt = option.dict_to_nonedict(opt)
if LATENT_DISTRIBUTION in NON_ARBITRARY_Z_INPUTS:
LATENT_CHANNEL_NUM = None
else:
TEST_IMAGE = None
# print to file and std_out simultaneously
sys.stdout = PrintLogger(opt['path']['log'])
print('\n**********' + util.get_timestamp() + '**********')
# Create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()):
assert dataset_opt['dataroot_LR'] is None or dataset_opt[
'dataroot_HR'] is None, 'Should not rely on saved LR versions when HR images are available. Downscaling images myself using CEM_imresize in the get_item routine.'
test_set = create_dataset(
dataset_opt,
specific_image=TEST_IMAGE,
kernel=None if opt['test'] is None else opt['test']['kernel'])
test_loader = create_dataloader(test_set, dataset_opt)
print('Number of test images in [{:s}]: {:d}'.format(
dataset_opt['name'], len(test_set)))
test_loaders.append(test_loader)
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option JSON file.')
parser.add_argument('-single_GPU',
action='store_true',
help='Utilize only one GPU')
if parser.parse_args().single_GPU:
available_GPUs = util.Assign_GPU()
else:
available_GPUs = util.Assign_GPU(max_GPUs=None)
opt = option.parse(parser.parse_args().opt,
is_train=True,
batch_size_multiplier=len(available_GPUs))
if not opt['train']['resume']:
util.mkdir_and_rename(
opt['path']
['experiments_root']) # Modify experiment name if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and \
not key == 'pretrained_model_G' and not key == 'pretrained_model_D'))
option.save(opt)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# print to file and std_out simultaneously
sys.stdout = PrintLogger(opt['path']['log'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
print("Random Seed: ", seed)
random.seed(seed)
torch.manual_seed(seed)
# create train and val dataloader
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
max_accumulation_steps = max([
opt['train']['grad_accumulation_steps_G'],
opt['train']['grad_accumulation_steps_D']
])
train_set = create_dataset(dataset_opt)
train_size = int(
math.ceil(len(train_set) / dataset_opt['batch_size']))
print('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'] *
max_accumulation_steps) #-current_step
total_epoches = int(math.ceil(total_iters / train_size))
print('Total epoches needed: {:d} for iters {:,d}'.format(
total_epoches, total_iters))
train_loader = create_dataloader(train_set, dataset_opt)
elif phase == 'val':
val_dataset_opt = dataset_opt
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
print('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
if max_accumulation_steps != 1:
model = create_model(opt, max_accumulation_steps)
else:
model = create_model(opt)
# create logger
logger = Logger(opt)
# Save validation set results as image collage:
SAVE_IMAGE_COLLAGE = True
per_image_saved_patch = min(
[min(im['HR'].shape[1:]) for im in val_loader.dataset]) - 2
num_val_images = len(val_loader.dataset)
val_images_collage_rows = int(np.floor(np.sqrt(num_val_images)))
while val_images_collage_rows > 1:
if np.round(num_val_images / val_images_collage_rows
) == num_val_images / val_images_collage_rows:
break
val_images_collage_rows -= 1
start_time = time.time()
min_accumulation_steps = min([
opt['train']['grad_accumulation_steps_G'],
opt['train']['grad_accumulation_steps_D']
])
save_GT_HR = True
lr_too_low = False
print('---------- Start training -------------')
last_saving_time = time.time()
recently_saved_models = deque(maxlen=4)
for epoch in range(int(math.floor(model.step / train_size)),
total_epoches):
for i, train_data in enumerate(train_loader):
gradient_step_num = model.step // max_accumulation_steps
not_within_batch = model.step % max_accumulation_steps == (
max_accumulation_steps - 1)
saving_step = (
(time.time() - last_saving_time) > 60 *
opt['logger']['save_checkpoint_freq']) and not_within_batch
if saving_step:
last_saving_time = time.time()
# save models
if lr_too_low or saving_step:
recently_saved_models.append(model.save(gradient_step_num))
model.save_log()
if len(recently_saved_models) > 3:
model_2_delete = recently_saved_models.popleft()
os.remove(model_2_delete)
if model.D_exists:
os.remove(model_2_delete.replace('_G.', '_D.'))
print('{}: Saving the model before iter {:d}.'.format(
datetime.now().strftime('%H:%M:%S'), gradient_step_num))
if lr_too_low:
break
if model.step > total_iters:
break
# training
model.feed_data(train_data)
model.optimize_parameters()
if not model.D_exists: #Avoid using the naive MultiLR scheduler when using adversarial loss
for scheduler in model.schedulers:
scheduler.step(model.gradient_step_num)
time_elapsed = time.time() - start_time
if not_within_batch: start_time = time.time()
# log
if gradient_step_num % opt['logger'][
'print_freq'] == 0 and not_within_batch:
logs = model.get_current_log()
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = gradient_step_num
print_rlt['time'] = time_elapsed
for k, v in logs.items():
print_rlt[k] = v
print_rlt['lr'] = model.get_current_learning_rate()
logger.print_format_results('train',
print_rlt,
keys_ignore_list=IGNORED_KEYS_LIST)
model.display_log_figure()
# validation
if not_within_batch and (gradient_step_num) % opt['train'][
'val_freq'] == 0: # and gradient_step_num>=opt['train']['D_init_iters']:
print_rlt = OrderedDict()
if model.generator_changed:
print('---------- validation -------------')
start_time = time.time()
if False and SAVE_IMAGE_COLLAGE and model.gradient_step_num % opt[
'train'][
'val_save_freq'] == 0: #Saving training images:
GT_image_collage = []
cur_train_results = model.get_current_visuals(
entire_batch=True)
train_psnrs = [
util.calculate_psnr(
util.tensor2img(
cur_train_results['SR'][im_num],
out_type=np.float32) * 255,
util.tensor2img(
cur_train_results['HR'][im_num],
out_type=np.float32) * 255)
for im_num in range(len(cur_train_results['SR']))
]
#Save latest training batch output:
save_img_path = os.path.join(
os.path.join(opt['path']['val_images']),
'{:d}_Tr_PSNR{:.3f}.png'.format(
gradient_step_num, np.mean(train_psnrs)))
util.save_img(
np.clip(
np.concatenate(
(np.concatenate([
util.tensor2img(
cur_train_results['HR'][im_num],
out_type=np.float32) * 255
for im_num in range(
len(cur_train_results['SR']))
], 0),
np.concatenate([
util.tensor2img(
cur_train_results['SR'][im_num],
out_type=np.float32) * 255
for im_num in range(
len(cur_train_results['SR']))
], 0)), 1), 0, 255).astype(np.uint8),
save_img_path)
Z_latent = [0] + ([-1, 1] if
opt['network_G']['latent_input'] else [])
print_rlt['psnr'] = 0
for cur_Z in Z_latent:
sr_images = model.perform_validation(
data_loader=val_loader,
cur_Z=cur_Z,
print_rlt=print_rlt,
save_GT_HR=save_GT_HR,
save_images=((model.gradient_step_num) %
opt['train']['val_save_freq'] == 0)
or save_GT_HR)
if logger.use_tb_logger:
logger.tb_logger.log_images(
'validation_Z%.2f' % (cur_Z),
[im[:, :, [2, 1, 0]] for im in sr_images],
model.gradient_step_num)
if save_GT_HR: # Save GT Uncomp images
save_GT_HR = False
model.log_dict['psnr_val'].append(
(gradient_step_num, print_rlt['psnr'] / len(Z_latent)))
else:
print('Skipping validation because generator is unchanged')
time_elapsed = time.time() - start_time
# Save to log
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = gradient_step_num
print_rlt['time'] = time_elapsed
model.display_log_figure()
logger.print_format_results('val',
print_rlt,
keys_ignore_list=IGNORED_KEYS_LIST)
print('-----------------------------------')
# update learning rate
if not_within_batch:
lr_too_low = model.update_learning_rate(gradient_step_num)
if lr_too_low:
print('Stopping training because LR is too low')
break
print('Saving the final model.')
model.save(gradient_step_num)
print('End of training.')
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, required=True, help='Path to option JSON file.')
parser.add_argument('-single_GPU', action='store_true',help='Utilize only one GPU')
parser.add_argument('-chroma', action='store_true',help='Training the chroma-channels generator')
if parser.parse_args().single_GPU:
available_GPUs = util.Assign_GPU(maxMemory=0.66)
else:
# available_GPUs = util.Assign_GPU(max_GPUs=None,maxMemory=0.8,maxLoad=0.8)
available_GPUs = util.Assign_GPU(max_GPUs=None)
opt = option.parse(parser.parse_args().opt, is_train=True,batch_size_multiplier=len(available_GPUs),name='JPEG'+('_chroma' if parser.parse_args().chroma else ''))
if not opt['train']['resume']:
util.mkdir_and_rename(opt['path']['experiments_root']) # Modify experiment name if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and \
not key == 'pretrained_model_G' and not key == 'pretrained_model_D'))
option.save(opt)
opt = option.dict_to_nonedict(opt) # Convert to NoneDict, which return None for missing key.
# print to file and std_out simultaneously
sys.stdout = PrintLogger(opt['path']['log'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
print("Random Seed: ", seed)
random.seed(seed)
torch.manual_seed(seed)
# create train and val dataloader
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
max_accumulation_steps = max([opt['train']['grad_accumulation_steps_G'], opt['train']['grad_accumulation_steps_D']])
train_set = create_dataset(dataset_opt)
train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size']))
print('Number of train images: {:,d}, iters: {:,d}'.format(len(train_set), train_size))
total_iters = int(opt['train']['niter']*max_accumulation_steps)#-current_step
total_epoches = int(math.ceil(total_iters / train_size))
print('Total epoches needed: {:d} for iters {:,d}'.format(total_epoches, total_iters))
train_loader = create_dataloader(train_set, dataset_opt)
elif phase == 'val':
val_dataset_opt = dataset_opt
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
print('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
DEBUG = False
# Create model
if DEBUG:
from models.base_model import BaseModel
model = BaseModel
model.step = 0
else:
model = create_model(opt,max_accumulation_steps,chroma_mode=opt['name'][:len('JPEG/chroma')]=='JPEG/chroma')
# create logger
logger = Logger(opt)
# Save validation set results as image collage:
SAVE_IMAGE_COLLAGE = True
start_time,start_time_gradient_step = time.time(),model.step // max_accumulation_steps
save_GT_Uncomp = True
lr_too_low = False
print('---------- Start training -------------')
last_saving_time = time.time()
recently_saved_models = deque(maxlen=4)
for epoch in range(int(math.floor(model.step / train_size)),total_epoches):
for i, train_data in enumerate(train_loader):
model.gradient_step_num = model.step // max_accumulation_steps
not_within_batch = model.step % max_accumulation_steps == (max_accumulation_steps - 1)
saving_step = ((time.time()-last_saving_time)>60*opt['logger']['save_checkpoint_freq']) and not_within_batch
if saving_step:
last_saving_time = time.time()
# save models
if lr_too_low or saving_step:
model.save_log()
recently_saved_models.append(model.save(model.gradient_step_num))
if len(recently_saved_models)>3:
model_2_delete = recently_saved_models.popleft()
os.remove(model_2_delete)
if model.D_exists:
os.remove(model_2_delete.replace('_G.','_D.'))
print('{}: Saving the model before iter {:d}.'.format(datetime.now().strftime('%H:%M:%S'),model.gradient_step_num))
if lr_too_low:
break
if model.step > total_iters:
break
# time_elapsed = time.time() - start_time
# if not_within_batch: start_time = time.time()
# log
if model.gradient_step_num % opt['logger']['print_freq'] == 0 and not_within_batch:
logs = model.get_current_log()
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = model.gradient_step_num
# time_elapsed = time.time() - start_time
print_rlt['time'] = (time.time() - start_time)/np.maximum(1,model.gradient_step_num-start_time_gradient_step)
start_time, start_time_gradient_step = time.time(), model.gradient_step_num
for k, v in logs.items():
print_rlt[k] = v
print_rlt['lr'] = model.get_current_learning_rate()
logger.print_format_results('train', print_rlt,keys_ignore_list=['avg_est_err'])
model.display_log_figure()
# validation
if (not_within_batch or i==0) and (model.gradient_step_num) % opt['train']['val_freq'] == 0: # and model.gradient_step_num>=opt['train']['D_init_iters']:
print_rlt = OrderedDict()
if model.generator_changed:
print('---------- validation -------------')
start_time = time.time()
if False and SAVE_IMAGE_COLLAGE and model.gradient_step_num%opt['train']['val_save_freq'] == 0: #Saving training images:
# GT_image_collage,quantized_image_collage = [],[]
cur_train_results = model.get_current_visuals(entire_batch=True)
train_psnrs = [util.calculate_psnr(util.tensor2img(cur_train_results['Decomp'][im_num], out_type=np.uint8,min_max=[0,255]),
util.tensor2img(cur_train_results['Uncomp'][im_num], out_type=np.uint8,min_max=[0,255])) for im_num in range(len(cur_train_results['Decomp']))]
#Save latest training batch output:
save_img_path = os.path.join(os.path.join(opt['path']['val_images']),
'{:d}_Tr_PSNR{:.3f}.png'.format(model.gradient_step_num, np.mean(train_psnrs)))
util.save_img(np.clip(np.concatenate((np.concatenate([util.tensor2img(cur_train_results['Uncomp'][im_num], out_type=np.uint8,min_max=[0,255]) for im_num in
range(len(cur_train_results['Decomp']))],0), np.concatenate(
[util.tensor2img(cur_train_results['Decomp'][im_num], out_type=np.uint8,min_max=[0,255]) for im_num in range(len(cur_train_results['Decomp']))],
0)), 1), 0, 255).astype(np.uint8), save_img_path)
Z_latent = [0]+([-0.5,0.5] if opt['network_G']['latent_input'] else [])
print_rlt['psnr'] = 0
for cur_Z in Z_latent:
model.perform_validation(data_loader=val_loader,cur_Z=cur_Z,print_rlt=print_rlt,GT_and_quantized=save_GT_Uncomp,
save_images=((model.gradient_step_num) % opt['train']['val_save_freq'] == 0) or save_GT_Uncomp)
if save_GT_Uncomp: # Save GT Uncomp images
save_GT_Uncomp = False
print_rlt['psnr'] /= len(Z_latent)
model.log_dict['psnr_val'].append((model.gradient_step_num,print_rlt['psnr']))
else:
print('Skipping validation because generator is unchanged')
# time_elapsed = time.time() - start_time
# Save to log
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = model.gradient_step_num
# print_rlt['time'] = time_elapsed
print_rlt['time'] = (time.time() - start_time)/np.maximum(1,model.gradient_step_num-start_time_gradient_step)
# model.display_log_figure()
# model.generator_changed = False
logger.print_format_results('val', print_rlt,keys_ignore_list=['avg_est_err'])
print('-----------------------------------')
model.feed_data(train_data,mixed_Y=True)
model.optimize_parameters()
# update learning rate
if not_within_batch:
lr_too_low = model.update_learning_rate(model.gradient_step_num)
# current_step += 1
if lr_too_low:
print('Stopping training because LR is too low')
break
print('Saving the final model.')
model.save(model.gradient_step_num)
model.save_log()
print('End of training.')
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=True)
util.mkdir_and_rename(
opt['path']['experiments_root']) # rename old experiments if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and \
not key == 'pretrain_model_G' and not key == 'pretrain_model_D'))
option.save(opt)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# print to file and std_out simultaneously
sys.stdout = PrintLogger(opt['path']['log'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
print("Random Seed: ", seed)
random.seed(seed)
torch.manual_seed(seed)
# create train and val dataloader
val_loaders = []
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']))
print('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epoches = int(math.ceil(total_iters / train_size))
print('Total epoches needed: {:d} for iters {:,d}'.format(
total_epoches, total_iters))
train_loader = create_dataloader(train_set, dataset_opt)
elif 'val' in phase:
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
print('Number of val images in [%s]: %d' %
(dataset_opt['name'], len(val_set)))
val_loaders.append(val_loader)
else:
raise NotImplementedError("Phase [%s] is not recognized." % phase)
assert train_loader is not None
# Create model
model = create_model(opt)
# create logger
logger = Logger(opt)
current_step = 0
start_time = time.time()
print('---------- Start training -------------')
for epoch in range(total_epoches):
for i, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
# level = random.randint(0, 80)
level = random.uniform(0, 80)
# train_data = add_Gaussian_noise(train_data, level)
train_data = add_dependent_noise(train_data, level)
model.feed_data(train_data)
# training
model.optimize_parameters(current_step)
time_elapsed = time.time() - start_time
start_time = time.time()
# log
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
print_rlt['time'] = time_elapsed
for k, v in logs.items():
print_rlt[k] = v
print_rlt['lr'] = model.get_current_learning_rate()
logger.print_format_results('train', print_rlt)
# save models
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
print('Saving the model at the end of iter {:d}.'.format(
current_step))
model.save(current_step)
# validation
if current_step % opt['train']['val_freq'] == 0:
print('---------- validation -------------')
for val_loader in val_loaders:
val_set_name = val_loader.dataset.opt['name']
print('validation [%s]...' % val_set_name)
print_rlt = OrderedDict()
start_time = time.time()
avg_l2 = 0.0
idx = 0
sigma_gt = int(val_set_name[-2:])
sigma_pre = []
for val_data in val_loader:
idx += 1
val_data = add_dependent_noise(val_data, sigma_gt)
model.feed_data(val_data, need_HR=False)
sigma = model.test_sigma().squeeze().float().cpu(
).item()
sigma_pre.append(sigma)
avg_l2 += (sigma - sigma_gt)**2
print('sigma: {}'.format(sigma_pre))
# log the sigma, time for each quality
time_elapsed = time.time() - start_time
log_val_name = 'l2_noise{}'.format(sigma_gt)
print_rlt[log_val_name] = avg_l2
print_rlt['time'] = time_elapsed
# Save to log
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
logger.print_format_results('val', print_rlt)
print('-----------------------------------')
# end of the validation
# update learning rate
model.update_learning_rate()
print('Saving the final model.')
model.save('latest')
print('End of training.')
def main():
# Settings
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, required=True, help='Path to option JSON file.')
opt = option.parse(parser.parse_args().opt) #load settings and initialize settings
util.mkdir_and_rename(opt['path']['experiments_root']) # rename old experiments if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and \
not key == 'saved_model'))
option.save(opt)
opt = option.dict_to_nonedict(opt) # Convert to NoneDict, which return None for missing key.
# Redirect all writes to the "txt" file
sys.stdout = PrintLogger(opt['path']['log'])
# create train and val dataloader
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']))
print('Number of train images: {:,d}, iters: {:,d}'.format(len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epoches = int(math.ceil(total_iters / train_size))
print('Total epoches needed: {:d} for iters {:,d}'.format(total_epoches, total_iters))
train_loader = create_dataloader(train_set, dataset_opt)
elif phase == 'val':
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
print('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 = create_model(opt)
# Create logger
logger = Logger(opt)
current_step = 0
start_time = time.time()
print('---------- Start training -------------')
for epoch in range(total_epoches):
for i, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
# training
model.feed_data(train_data)
model.optimize_parameters(current_step)
time_elapsed = time.time() - start_time
start_time = time.time()
# log
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
print_rlt['time'] = time_elapsed
for k, v in logs.items():
print_rlt[k] = v
print_rlt['lr'] = model.get_current_learning_rate()
logger.print_format_results('train', print_rlt)
# save models
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
print('Saving the model at the end of iter {:d}.'.format(current_step))
model.save(current_step)
# validation
if current_step % opt['train']['val_freq'] == 0:
print('---------- validation -------------')
start_time = time.time()
avg_psnr = 0.0
avg_ssim =0.0
idx = 0
for val_data in val_loader:
idx += 1
img_name = os.path.splitext(os.path.basename(val_data['GT_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()
out_img = util.tensor2img(visuals['Output'])
gt_img = util.tensor2img(visuals['ground_truth']) # uint8
# Save output images for reference
save_img_path = os.path.join(img_dir, '{:s}_{:d}.png'.format(\
img_name, current_step))
util.save_img(out_img, save_img_path)
# calculate PSNR
if len(gt_img.shape) == 2:
gt_img = np.expand_dims(gt_img, axis=2)
out_img = np.expand_dims(out_img, axis=2)
crop_border = opt['scale']
cropped_out_img = out_img[crop_border:-crop_border, crop_border:-crop_border, :]
cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :]
if gt_img.shape[2] == 3: # RGB image
cropped_out_img_y = bgr2ycbcr(cropped_out_img, only_y=True)
cropped_gt_img_y = bgr2ycbcr(cropped_gt_img, only_y=True)
avg_psnr += util.psnr(cropped_out_img_y, cropped_gt_img_y)
avg_ssim += util.ssim(cropped_out_img_y, cropped_gt_img_y, multichannel=False)
else:
avg_psnr += util.psnr(cropped_out_img, cropped_gt_img)
avg_ssim += util.ssim(cropped_out_img, cropped_gt_img, multichannel=True)
avg_psnr = avg_psnr / idx
avg_ssim = avg_ssim / idx
time_elapsed = time.time() - start_time
# Save to log
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
print_rlt['time'] = time_elapsed
print_rlt['psnr'] = avg_psnr
print_rlt['ssim'] = avg_ssim
logger.print_format_results('val', print_rlt)
print('-----------------------------------')
# update learning rate
model.update_learning_rate()
print('Saving the final model.')
model.save('latest')
print('End of training.')
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=True)
util.mkdir_and_rename(
opt['path']['experiments_root']) # rename old experiments if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and \
not key == 'pretrain_model_G' and not key == 'pretrain_model_D'))
option.save(opt)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# print to file and std_out simultaneously
sys.stdout = PrintLogger(opt['path']['log'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
print("Random Seed: ", seed)
random.seed(seed)
torch.manual_seed(seed)
# create train and val dataloader
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']))
print('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epoches = int(math.ceil(total_iters / train_size))
print('Total epoches needed: {:d} for iters {:,d}'.format(
total_epoches, total_iters))
train_loader = create_dataloader(train_set, dataset_opt)
batch_size_per_month = dataset_opt['batch_size']
batch_size_per_day = int(
opt['datasets']['train']['batch_size_per_day'])
num_month = int(opt['train']['num_month'])
num_day = int(opt['train']['num_day'])
use_dci = false if 'use_dci' not in opt['train'] else opt['train'][
'use_dci']
elif phase == 'val':
val_dataset_opt = dataset_opt
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
print('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 = create_model(opt)
# create logger
logger = Logger(opt)
current_step = 0
start_time = time.time()
print('---------- Start training -------------')
for epoch in range(num_month):
for i, train_data in enumerate(train_loader):
# get the code
if use_dci:
cur_month_code = get_code_for_data(model, train_data, opt)
else:
cur_month_code = get_code(model, train_data, opt)
for j in range(num_day):
current_step += 1
if current_step > total_iters:
break
# get the sliced data
cur_day_batch_start_idx = (
j * batch_size_per_day) % batch_size_per_month
cur_day_batch_end_idx = cur_day_batch_start_idx + batch_size_per_day
if cur_day_batch_end_idx > batch_size_per_month:
cur_day_batch_idx = np.hstack(
(np.arange(cur_day_batch_start_idx,
batch_size_per_month),
np.arange(cur_day_batch_end_idx -
batch_size_per_month)))
else:
cur_day_batch_idx = slice(cur_day_batch_start_idx,
cur_day_batch_end_idx)
cur_day_train_data = {
'LR': train_data['LR'][cur_day_batch_idx],
'HR': train_data['HR'][cur_day_batch_idx]
}
code = cur_month_code[cur_day_batch_idx]
# training
model.feed_data(cur_day_train_data, code=code)
model.optimize_parameters(current_step)
time_elapsed = time.time() - start_time
start_time = time.time()
# log
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
print_rlt['time'] = time_elapsed
for k, v in logs.items():
print_rlt[k] = v
print_rlt['lr'] = model.get_current_learning_rate()
logger.print_format_results('train', print_rlt)
# save models
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
print('Saving the model at the end of iter {:d}.'.format(
current_step))
model.save(current_step)
# validation
if current_step % opt['train']['val_freq'] == 0:
print('---------- validation -------------')
start_time = time.time()
avg_psnr = 0.0
idx = 0
for val_data in val_loader:
idx += 1
img_name = os.path.splitext(
os.path.basename(val_data['LR_path'][0]))[0]
img_dir = os.path.join(opt['path']['val_images'],
img_name)
util.mkdir(img_dir)
if 'zero_code' in opt['train'] and opt['train'][
'zero_code']:
code_val = torch.zeros(
val_data['LR'].shape[0],
int(opt['network_G']['in_code_nc']),
val_data['LR'].shape[2],
val_data['LR'].shape[3])
elif 'rand_code' in opt['train'] and opt['train'][
'rand_code']:
code_val = torch.rand(
val_data['LR'].shape[0],
int(opt['network_G']['in_code_nc']),
val_data['LR'].shape[2],
val_data['LR'].shape[3])
else:
code_val = torch.randn(
val_data['LR'].shape[0],
int(opt['network_G']['in_code_nc']),
val_data['LR'].shape[2],
val_data['LR'].shape[3])
model.feed_data(val_data, code=code_val)
model.test()
visuals = model.get_current_visuals()
sr_img = util.tensor2img(visuals['SR']) # uint8
gt_img = util.tensor2img(visuals['HR']) # uint8
# Save SR images for reference
run_index = opt['name'].split("_")[2]
save_img_path = os.path.join(img_dir, 'srim_{:s}_{:s}_{:d}.png'.format( \
run_index, img_name, current_step))
util.save_img(sr_img, save_img_path)
# calculate PSNR
crop_size = opt['scale']
cropped_sr_img = sr_img[crop_size:-crop_size,
crop_size:-crop_size, :]
cropped_gt_img = gt_img[crop_size:-crop_size,
crop_size:-crop_size, :]
avg_psnr += util.psnr(cropped_sr_img, cropped_gt_img)
avg_psnr = avg_psnr / idx
time_elapsed = time.time() - start_time
# Save to log
print_rlt = OrderedDict()
print_rlt['model'] = opt['model']
print_rlt['epoch'] = epoch
print_rlt['iters'] = current_step
print_rlt['time'] = time_elapsed
print_rlt['psnr'] = avg_psnr
logger.print_format_results('val', print_rlt)
print('-----------------------------------')
# update learning rate
model.update_learning_rate()
print('Saving the final model.')
model.save('latest')
print('End of training.')