def main(args):
# parse the options file
print(f'Parsing file {args.options_file}...')
opt = option.parse(args.options_file, is_train=False)
opt = option.dict_to_nonedict(opt)
print('Loading test images...')
test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()):
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set, dataset_opt)
test_loaders.append(test_loader)
model = create_model(opt)
generator = model.netG.module
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
print(f'Testing on dataset {test_set_name}')
psnr_vals, ssim_vals = evaluate_generator(generator, test_loader, opt)
psnr_val = np.mean(psnr_vals)
ssim_val = np.mean(ssim_vals)
print(
f'Mean PSNR and SSIM for {test_set_name} on original model are: [{psnr_val}, {ssim_val}]'
)
# The input shape is chosen arbitrarily to generate dummy input for creating quantsim object
input_shapes = (1, 3, 24, 24)
sim = quantsim.QuantizationSimModel(
generator,
input_shapes=input_shapes,
quant_scheme=args.quant_scheme,
default_output_bw=args.default_output_bw,
default_param_bw=args.default_param_bw)
evaluate_func = partial(evaluate_generator, options=opt)
sim.compute_encodings(evaluate_func, test_loaders[0])
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
print(f'Testing on dataset {test_set_name}')
psnr_vals, ssim_vals = evaluate_generator(sim.model,
test_loader,
opt,
output_dir=args.output_dir)
psnr_val = np.mean(psnr_vals)
ssim_val = np.mean(ssim_vals)
print(
f'Mean PSNR and SSIM for {test_set_name} on quantized model are: [{psnr_val}, {ssim_val}]'
)
def main():
# parse options, set distributed setting, set ramdom seed
opt = parse_options(is_train=False)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# mkdir and initialize loggers
make_exp_dirs(opt)
log_file = osp.join(opt['path']['log'],
f"test_{opt['name']}_{get_time_str()}.log")
logger = get_root_logger(logger_name='codes',
log_level=logging.INFO,
log_file=log_file)
logger.info(dict2str(opt))
# create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()):
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=None,
seed=opt['manual_seed'])
logger.info(
f"Number of test images in {dataset_opt['name']}: {len(test_set)}")
test_loaders.append(test_loader)
# create model
model = create_model(opt)
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
logger.info(f'Testing {test_set_name}...')
model.validation(test_loader,
current_iter=opt['name'],
tb_logger=None,
save_img=opt['val']['save_img'])
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, default='/home/daniel/storage/BasicSR/codes/options/train/train_cityscapeSFT.json',
help='Path to option JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=True)
opt = option.dict_to_nonedict(opt) # Convert to NoneDict, which return None for missing key.
print(opt)
# TODO: check if this folder setting code works
if opt['path']['resume_state']: # resuming training
resume_state = torch.load(opt['path']['resume_state'])
else: # training from scratch
resume_state = None
util.mkdir_and_rename(opt['path']['experiments_root']) # rename old 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))
# TODO: deal with the use of "logging" in this crackden
writer = SummaryWriter()
writer.add_text('Options', str(opt))
torch.backends.cudnn.benchmark = True
# create train and val dataloader
train_loader = cdl.build_custom_dataloader(opt, 'train')
val_loader = cdl.build_custom_dataloader(opt, 'val')
model = create_model(opt)
# training
current_step = 0
start_epoch = 0
total_epochs = 2000
epoch_len = len(train_loader)
print('starting training')
for epoch in range(start_epoch, total_epochs):
tic = time.time()
for _, train_data in enumerate(train_loader):
current_step += 1
# model.update_learning_rate()
model.feed_data(train_data)
model.optimize_parameters(current_step)
if current_step % 30 == 1:
losses = model.get_current_losses()
epoch_step = current_step % epoch_len
visuals = model.get_current_visuals()
gt_img = util.tensor2img(visuals['HR']) # uint8
sr_img = util.tensor2img(visuals['SR']) # uint8
writer.add_image('GT_train', gt_img, current_step, dataformats='HWC')
writer.add_image('SR_train', sr_img, current_step, dataformats='HWC')
print(f'step {epoch_step}/{epoch_len} || loss: {losses}')
writer.add_scalars('Losses', losses, current_step)
if current_step % opt['train']['val_freq'] == 0:
print('validating model')
avg_psnr = validate(current_step, model, opt, val_loader)
writer.add_scalar('Validation loss', avg_psnr, current_step)
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
print('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
# more logging
if epoch % 10 == 0:
print(f'epoch {epoch} done, took {time.time()-tic}s')
def main():
# parse options, set distributed setting, set ramdom seed
opt = parse_options(is_train=True)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# load resume states if necessary
if opt['path'].get('resume_state'):
device_id = torch.cuda.current_device()
resume_state = torch.load(
opt['path']['resume_state'],
map_location=lambda storage, loc: storage.cuda(device_id))
else:
resume_state = None
# mkdir for experiments and logger
if resume_state is None:
make_exp_dirs(opt)
if opt['logger'].get('use_tb_logger') and 'debug' not in opt[
'name'] and opt['rank'] == 0:
mkdir_and_rename(osp.join('tb_logger', opt['name']))
# initialize loggers
logger, tb_logger = init_loggers(opt)
# create train and validation dataloaders
result = create_train_val_dataloader(opt, logger)
train_loader, train_sampler, val_loader, total_epochs, total_iters = result
# create model
if resume_state: # resume training
check_resume(opt, resume_state['iter'])
model = create_model(opt)
model.resume_training(resume_state) # handle optimizers and schedulers
logger.info(f"Resuming training from epoch: {resume_state['epoch']}, "
f"iter: {resume_state['iter']}.")
start_epoch = resume_state['epoch']
current_iter = resume_state['iter']
else:
model = create_model(opt)
start_epoch = 0
current_iter = 0
# create message logger (formatted outputs)
msg_logger = MessageLogger(opt, current_iter, tb_logger)
# dataloader prefetcher
prefetch_mode = opt['datasets']['train'].get('prefetch_mode')
if prefetch_mode is None or prefetch_mode == 'cpu':
prefetcher = CPUPrefetcher(train_loader)
elif prefetch_mode == 'cuda':
prefetcher = CUDAPrefetcher(train_loader, opt)
logger.info(f'Use {prefetch_mode} prefetch dataloader')
if opt['datasets']['train'].get('pin_memory') is not True:
raise ValueError('Please set pin_memory=True for CUDAPrefetcher.')
else:
raise ValueError(f'Wrong prefetch_mode {prefetch_mode}.'
"Supported ones are: None, 'cuda', 'cpu'.")
# training
logger.info(
f'Start training from epoch: {start_epoch}, iter: {current_iter}')
data_time, iter_time = time.time(), time.time()
start_time = time.time()
for epoch in range(start_epoch, total_epochs + 1):
train_sampler.set_epoch(epoch)
prefetcher.reset()
train_data = prefetcher.next()
while train_data is not None:
data_time = time.time() - data_time
current_iter += 1
if current_iter > total_iters:
break
# update learning rate
model.update_learning_rate(current_iter,
warmup_iter=opt['train'].get(
'warmup_iter', -1))
# training
model.feed_data(train_data)
model.optimize_parameters(current_iter)
iter_time = time.time() - iter_time
# log
if current_iter % opt['logger']['print_freq'] == 0:
log_vars = {'epoch': epoch, 'iter': current_iter}
log_vars.update({'lrs': model.get_current_learning_rate()})
log_vars.update({'time': iter_time, 'data_time': data_time})
log_vars.update(model.get_current_log())
msg_logger(log_vars)
# save models and training states
if current_iter % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(epoch, current_iter)
# validation
if opt.get('val') is not None and (current_iter %
opt['val']['val_freq'] == 0):
model.validation(val_loader, current_iter, tb_logger,
opt['val']['save_img'])
data_time = time.time()
iter_time = time.time()
train_data = prefetcher.next()
# end of iter
# end of epoch
consumed_time = str(
datetime.timedelta(seconds=int(time.time() - start_time)))
logger.info(f'End of training. Time consumed: {consumed_time}')
logger.info('Save the latest model.')
model.save(epoch=-1, current_iter=-1) # -1 stands for the latest
if opt.get('val') is not None:
model.validation(val_loader, current_iter, tb_logger,
opt['val']['save_img'])
if tb_logger:
tb_logger.close()
# 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()):
test_set = create_dataset(dataset_opt)
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)
# Create model
model = create_model(opt)
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
print('\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'] = []
need_ground_truth = True