def test(opt):
# logging
logger = base_utils.get_logger('base')
if opt['verbose']:
logger.info('{} Configurations {}'.format('=' * 20, '=' * 20))
base_utils.print_options(opt, logger)
# infer and evaluate performance for each model
for load_path in opt['model']['generator']['load_path_lst']:
# setup model index
model_idx = osp.splitext(osp.split(load_path)[-1])[0]
# log
logger.info('=' * 40)
logger.info('Testing model: {}'.format(model_idx))
logger.info('=' * 40)
# create model
opt['model']['generator']['load_path'] = load_path
model = define_model(opt)
model_idx = osp.basename(opt['model']['generator']['load_path']).split('.')[0]
# for each test dataset
for dataset_idx in sorted(opt['dataset'].keys()):
# use dataset with prefix `test`
if not dataset_idx.startswith('test'):
continue
validate(opt, model, logger, dataset_idx, model_idx, compute_metrics=False)
logger.info('-' * 40)
# logging
logger.info('Finish testing')
logger.info('=' * 40)
def test(opt):
# logging
logger = base_utils.get_logger('base')
if opt['verbose']:
logger.info('{} Configurations {}'.format('=' * 20, '=' * 20))
base_utils.print_options(opt, logger)
# infer and evaluate performance for each model
for load_path in opt['model']['generator']['load_path_lst']:
# setup model index
model_idx = osp.splitext(osp.split(load_path)[-1])[0]
# log
logger.info('=' * 40)
logger.info('Testing model: {}'.format(model_idx))
logger.info('=' * 40)
# create model
opt['model']['generator']['load_path'] = load_path
model = define_model(opt)
# for each test dataset
for dataset_idx in sorted(opt['dataset'].keys()):
# use dataset with prefix `test`
if not dataset_idx.startswith('test'):
continue
ds_name = opt['dataset'][dataset_idx]['name']
logger.info('Testing on {}: {}'.format(dataset_idx, ds_name))
# create data loader
test_loader = create_dataloader(opt, dataset_idx=dataset_idx)
# infer and store results for each sequence
for i, data in enumerate(test_loader):
# fetch data
lr_data = data['lr'][0]
seq_idx = data['seq_idx'][0]
frm_idx = [frm_idx[0] for frm_idx in data['frm_idx']]
# infer
hr_seq = model.infer(lr_data) # thwc|rgb|uint8
# save results (optional)
if opt['test']['save_res']:
res_dir = osp.join(opt['test']['res_dir'], ds_name,
model_idx)
res_seq_dir = osp.join(res_dir, seq_idx)
data_utils.save_sequence(res_seq_dir,
hr_seq,
frm_idx,
to_bgr=True)
logger.info('-' * 40)
# logging
logger.info('Finish testing')
logger.info('=' * 40)
def __init__(self, opt):
self.opt = opt
self.verbose = opt['verbose']
self.scale = opt['scale']
self.logger = get_logger('base')
self.device = torch.device(opt['device'])
self.is_train = opt['is_train']
if self.is_train:
self.ckpt_dir = opt['train']['ckpt_dir']
self.log_decay = opt['logger'].get('decay', 0.99)
self.log_dict = OrderedDict()
self.running_log_dict = OrderedDict()
def profile(opt, lr_size, test_speed=False):
# logging
logger = base_utils.get_logger('base')
logger.info('{} Model Information {}'.format('='*20, '='*20))
base_utils.print_options(opt['model']['generator'], logger)
# basic configs
scale = opt['scale']
device = torch.device(opt['device'])
# create model
net_G = define_generator(opt).to(device)
# get dummy input
dummy_input_dict = net_G.generate_dummy_input(lr_size)
for key in dummy_input_dict.keys():
dummy_input_dict[key] = dummy_input_dict[key].to(device)
# profile
register(net_G, dummy_input_dict)
gflops, params = profile_model(net_G)
logger.info('-' * 40)
logger.info('Super-resolute data from {}x{}x{} to {}x{}x{}'.format(
*lr_size, lr_size[0], lr_size[1]*scale, lr_size[2]*scale))
logger.info('Parameters (x10^6): {:.3f}'.format(params/1e6))
logger.info('FLOPs (x10^9): {:.3f}'.format(gflops))
logger.info('-' * 40)
# test running speed
if test_speed:
n_test = 3
tot_time = 0
for i in range(n_test):
start_time = time.time()
with torch.no_grad():
_ = net_G(**dummy_input_dict)
end_time = time.time()
tot_time += end_time - start_time
logger.info('Speed (FPS): {:.3f} (averaged for {} runs)'.format(
n_test / tot_time, n_test))
logger.info('-' * 40)
def display_results(self):
logger = base_utils.get_logger('base')
# per sequence results
for seq, metric_dict_per_seq in self.metric_dict.items():
logger.info('Sequence: {}'.format(seq))
for metric_type in self.metric_opt.keys():
mult = getattr(self, '{}_mult'.format(metric_type.lower()))
logger.info('\t{}: {:.6f} (x{})'.format(
metric_type,
mult * np.mean(metric_dict_per_seq[metric_type]), mult))
# average results
logger.info('Average')
metric_avg_dict = self.get_averaged_results()
for metric_type, avg_result in metric_avg_dict.items():
mult = getattr(self, '{}_mult'.format(metric_type.lower()))
logger.info('\t{}: {:.6f} (x{})'.format(metric_type,
mult * avg_result, mult))
def train(opt):
# logging
logger = base_utils.get_logger('base')
logger.info('{} Options {}'.format('='*20, '='*20))
base_utils.print_options(opt, logger)
# create data loader
train_loader = create_dataloader(opt, dataset_idx='train')
# create downsampling kernels for BD degradation
kernel = data_utils.create_kernel(opt)
# create model
model = define_model(opt)
# training configs
total_sample = len(train_loader.dataset)
iter_per_epoch = len(train_loader)
total_iter = opt['train']['total_iter']
total_epoch = int(math.ceil(total_iter / iter_per_epoch))
start_iter, iter = opt['train']['start_iter'], 0
test_freq = opt['test']['test_freq']
log_freq = opt['logger']['log_freq']
ckpt_freq = opt['logger']['ckpt_freq']
logger.info('Number of training samples: {}'.format(total_sample))
logger.info('Total epochs needed: {} for {} iterations'.format(
total_epoch, total_iter))
# train
for epoch in range(total_epoch):
for data in train_loader:
# update iter
iter += 1
curr_iter = start_iter + iter
if iter > total_iter:
logger.info('Finish training')
break
# update learning rate
model.update_learning_rate()
# prepare data
data = prepare_data(opt, data, kernel)
# train for a mini-batch
model.train(data)
# update running log
model.update_running_log()
# log
if log_freq > 0 and iter % log_freq == 0:
# basic info
msg = '[epoch: {} | iter: {}'.format(epoch, curr_iter)
for lr_type, lr in model.get_current_learning_rate().items():
msg += ' | {}: {:.2e}'.format(lr_type, lr)
msg += '] '
# loss info
log_dict = model.get_running_log()
msg += ', '.join([
'{}: {:.3e}'.format(k, v) for k, v in log_dict.items()])
logger.info(msg)
# save model
if ckpt_freq > 0 and iter % ckpt_freq == 0:
model.save(curr_iter)
# evaluate performance
if test_freq > 0 and iter % test_freq == 0:
# setup model index
model_idx = 'G_iter{}'.format(curr_iter)
# for each testset
for dataset_idx in sorted(opt['dataset'].keys()):
# use dataset with prefix `test`
if not dataset_idx.startswith('test'):
continue
ds_name = opt['dataset'][dataset_idx]['name']
logger.info(
'Testing on {}: {}'.format(dataset_idx, ds_name))
# create data loader
test_loader = create_dataloader(opt, dataset_idx=dataset_idx)
# define metric calculator
metric_calculator = MetricCalculator(opt)
# infer and compute metrics for each sequence
for data in test_loader:
# fetch data
lr_data = data['lr'][0]
seq_idx = data['seq_idx'][0]
frm_idx = [frm_idx[0] for frm_idx in data['frm_idx']]
# infer
hr_seq = model.infer(lr_data) # thwc|rgb|uint8
# save results (optional)
if opt['test']['save_res']:
res_dir = osp.join(
opt['test']['res_dir'], ds_name, model_idx)
res_seq_dir = osp.join(res_dir, seq_idx)
data_utils.save_sequence(
res_seq_dir, hr_seq, frm_idx, to_bgr=True)
# compute metrics for the current sequence
true_seq_dir = osp.join(
opt['dataset'][dataset_idx]['gt_seq_dir'], seq_idx)
metric_calculator.compute_sequence_metrics(
seq_idx, true_seq_dir, '', pred_seq=hr_seq)
# save/print metrics
if opt['test'].get('save_json'):
# save results to json file
json_path = osp.join(
opt['test']['json_dir'], '{}_avg.json'.format(ds_name))
metric_calculator.save_results(
model_idx, json_path, override=True)
else:
# print directly
metric_calculator.display_results()
def test(opt):
# logging
logger = base_utils.get_logger('base')
if opt['verbose']:
logger.info('{} Configurations {}'.format('=' * 20, '=' * 20))
base_utils.print_options(opt, logger)
# infer and evaluate performance for each model
for load_path in opt['model']['generator']['load_path_lst']:
# setup model index
model_idx = osp.splitext(osp.split(load_path)[-1])[0]
# log
logger.info('=' * 40)
logger.info('Testing model: {}'.format(model_idx))
logger.info('=' * 40)
# create model
opt['model']['generator']['load_path'] = load_path
model = define_model(opt)
# for each test dataset
for dataset_idx in sorted(opt['dataset'].keys()):
# use dataset with prefix `test`
if not dataset_idx.startswith('test'):
continue
ds_name = opt['dataset'][dataset_idx]['name']
logger.info('Testing on {}: {}'.format(dataset_idx, ds_name))
# define metric calculator
try:
metric_calculator = MetricCalculator(opt)
except:
print('No metirc need to compute!')
# create data loader
test_loader = create_dataloader(opt, dataset_idx=dataset_idx)
# infer and store results for each sequence
for i, data in enumerate(test_loader):
# fetch data
lr_data = data['lr'][0]
seq_idx = data['seq_idx'][0]
frm_idx = [frm_idx[0] for frm_idx in data['frm_idx']]
# infer
hr_seq = model.infer(lr_data) # thwc|rgb|uint8
# save results (optional)
if opt['test']['save_res']:
res_dir = osp.join(opt['test']['res_dir'], ds_name,
model_idx)
res_seq_dir = osp.join(res_dir, seq_idx)
data_utils.save_sequence(res_seq_dir,
hr_seq,
frm_idx,
to_bgr=True)
# compute metrics for the current sequence
true_seq_dir = osp.join(
opt['dataset'][dataset_idx]['gt_seq_dir'], seq_idx)
try:
metric_calculator.compute_sequence_metrics(seq_idx,
true_seq_dir,
'',
pred_seq=hr_seq)
except:
print('No metirc need to compute!')
# save/print metrics
try:
if opt['test'].get('save_json'):
# save results to json file
json_path = osp.join(opt['test']['json_dir'],
'{}_avg.json'.format(ds_name))
metric_calculator.save_results(model_idx,
json_path,
override=True)
else:
# print directly
metric_calculator.display_results()
except:
print('No metirc need to save!')
logger.info('-' * 40)
# logging
logger.info('Finish testing')
logger.info('=' * 40)
def train(opt):
# logging
logger = base_utils.get_logger('base')
logger.info('{} Options {}'.format('='*20, '='*20))
base_utils.print_options(opt, logger)
# create data loader
train_loader = create_dataloader(opt, dataset_idx='train')
# create downsampling kernels for BD degradation
kernel = data_utils.create_kernel(opt)
# create model
model = define_model(opt)
# training configs
total_sample = len(train_loader.dataset)
iter_per_epoch = len(train_loader)
total_iter = opt['train']['total_iter']
total_epoch = int(math.ceil(total_iter / iter_per_epoch))
curr_iter = opt['train']['start_iter']
test_freq = opt['test']['test_freq']
log_freq = opt['logger']['log_freq']
ckpt_freq = opt['logger']['ckpt_freq']
sigma_freq = opt['dataset']['degradation'].get('sigma_freq', 0)
sigma_inc = opt['dataset']['degradation'].get('sigma_inc', 0)
sigma_max = opt['dataset']['degradation'].get('sigma_max', 10)
logger.info('Number of training samples: {}'.format(total_sample))
logger.info('Total epochs needed: {} for {} iterations'.format(
total_epoch, total_iter))
print('device count:', torch.cuda.device_count())
# train
for epoch in range(total_epoch):
for data in tqdm(train_loader):
# update iter
curr_iter += 1
if curr_iter > total_iter:
logger.info('Finish training')
break
# update learning rate
model.update_learning_rate()
# prepare data
data = prepare_data(opt, data, kernel)
# train for a mini-batch
model.train(data)
# update running log
model.update_running_log()
# log
if log_freq > 0 and curr_iter % log_freq == 0:
# basic info
msg = '[epoch: {} | iter: {}'.format(epoch, curr_iter)
for lr_type, lr in model.get_current_learning_rate().items():
msg += ' | {}: {:.2e}'.format(lr_type, lr)
msg += '] '
# loss info
log_dict = model.get_running_log()
msg += ', '.join([
'{}: {:.3e}'.format(k, v) for k, v in log_dict.items()])
if opt['dataset']['degradation']['type'] == 'BD':
msg += ' | Sigma: {}'.format(opt['dataset']['degradation']['sigma'])
logger.info(msg)
# save model
if ckpt_freq > 0 and curr_iter % ckpt_freq == 0:
model.save(curr_iter)
# evaluate performance
if test_freq > 0 and curr_iter % test_freq == 0:
# setup model index
model_idx = 'G_iter{}'.format(curr_iter)
if opt['dataset']['degradation']['type'] == 'BD':
model_idx = model_idx + str(opt['dataset']['degradation']['sigma'])
# for each testset
for dataset_idx in sorted(opt['dataset'].keys()):
# use dataset with prefix `test`
if not dataset_idx.startswith('validate'):
continue
validate(opt, model, logger, dataset_idx, model_idx)
# schedule sigma
if opt['dataset']['degradation']['type'] == 'BD':
if sigma_freq > 0 and (epoch + 1) % sigma_freq == 0:
current_sigma = opt['dataset']['degradation']['sigma']
opt['dataset']['degradation']['sigma'] = min(current_sigma + sigma_inc, sigma_max)
kernel = data_utils.create_kernel(opt)
# __getitem__ in custom dataset class uses some crop that depends sigma
# it is crucial to change this cropsize accordingly if sigma is being changed
train_loader.dataset.change_cropsize(opt['dataset']['degradation']['sigma'])
print('kernel changed')
