def setup_logging(opt, resume_state, rank):
tb_logger = None
if rank <= 0: # normal training (rank -1) OR distributed training (rank 0)
if resume_state is None:
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))
# 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
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
else:
util.setup_logger('base',
opt['path']['log'],
'train',
level=logging.INFO,
screen=True)
logger = logging.getLogger('base')
return logger, tb_logger
def get_options(json_path):
"""options"""
# parser = argparse.ArgumentParser()
# parser.add_argument(
# '-opt', type=str, required=True, help='Path to options JSON file.')
# opt = option.parse(parser.parse_args().opt, is_train=False)
is_train = False
opt = option.parse(json_path, is_train)
util.mkdirs((path for key, path in opt['path'].items()
if not key == 'pretrain_model_G'))
opt = option.dict_to_nonedict(opt)
util.setup_logger(None,
opt['path']['log'],
'test',
level=logging.INFO,
screen=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt))
return opt, logger
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument("-opt", type=str, help="Path to option YAML 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:
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))
# 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)
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
tb_logger = SummaryWriter(log_dir="../tb_logger/" + opt["name"])
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)
# 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.benchmark = 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 = create_model(opt)
print("Model created!")
# 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)
for _, train_data in enumerate(train_loader):
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)
# log
if current_step % opt["logger"]["print_freq"] == 0:
logs = model.get_current_log()
message = "[epoch:{:3d}, iter:{:8,d}, lr:(".format(
epoch, current_step)
for v in model.get_current_learning_rate():
message += "{:.3e},".format(v)
message += ")] "
for k, v in logs.items():
message += "{:s}: {:.4e} ".format(k, v)
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
# validation
if opt["datasets"].get(
"val",
None) and current_step % opt["train"]["val_freq"] == 0:
# image restoration validation
if opt["model"] in ["sr", "srgan"] and rank <= 0:
# does not support multi-GPU validation
pbar = util.ProgressBar(len(val_loader))
avg_psnr = 0.0
idx = 0
for val_data in val_loader:
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["rlt"]) # uint8
gt_img = util.tensor2img(visuals["GT"]) # 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
sr_img, gt_img = util.crop_border([sr_img, gt_img],
opt["scale"])
avg_psnr += util.calculate_psnr(sr_img, gt_img)
pbar.update("Test {}".format(img_name))
avg_psnr = avg_psnr / idx
# log
logger.info("# Validation # PSNR: {:.4e}".format(avg_psnr))
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
tb_logger.add_scalar("psnr", avg_psnr, current_step)
else: # video restoration validation
if opt["dist"]:
# multi-GPU testing
psnr_rlt = {} # with border and center frames
if rank == 0:
pbar = util.ProgressBar(len(val_set))
for idx in range(rank, len(val_set), world_size):
val_data = val_set[idx]
val_data["LQs"].unsqueeze_(0)
val_data["GT"].unsqueeze_(0)
folder = val_data["folder"]
idx_d, max_idx = val_data["idx"].split("/")
idx_d, max_idx = int(idx_d), int(max_idx)
if psnr_rlt.get(folder, None) is None:
psnr_rlt[folder] = torch.zeros(
max_idx,
dtype=torch.float32,
device="cuda")
model.feed_data(val_data)
model.test()
visuals = model.get_current_visuals()
rlt_img = util.tensor2img(visuals["rlt"]) # uint8
gt_img = util.tensor2img(visuals["GT"]) # uint8
# calculate PSNR
psnr_rlt[folder][idx_d] = util.calculate_psnr(
rlt_img, gt_img)
if rank == 0:
for _ in range(world_size):
pbar.update("Test {} - {}/{}".format(
folder, idx_d, max_idx))
# collect data
for _, v in psnr_rlt.items():
dist.reduce(v, 0)
dist.barrier()
if rank == 0:
psnr_rlt_avg = {}
psnr_total_avg = 0.0
for k, v in psnr_rlt.items():
psnr_rlt_avg[k] = torch.mean(v).cpu().item()
psnr_total_avg += psnr_rlt_avg[k]
psnr_total_avg /= len(psnr_rlt)
log_s = "# Validation # PSNR: {:.4e}:".format(
psnr_total_avg)
for k, v in psnr_rlt_avg.items():
log_s += " {}: {:.4e}".format(k, v)
logger.info(log_s)
if opt["use_tb_logger"] and "debug" not in opt[
"name"]:
tb_logger.add_scalar("psnr_avg",
psnr_total_avg,
current_step)
for k, v in psnr_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
else:
pbar = util.ProgressBar(len(val_loader))
psnr_rlt = {} # with border and center frames
psnr_rlt_avg = {}
psnr_total_avg = 0.0
for val_data in val_loader:
folder = val_data["folder"][0]
idx_d, max_id = val_data["idx"][0].split("/")
# border = val_data['border'].item()
if psnr_rlt.get(folder, None) is None:
psnr_rlt[folder] = []
model.feed_data(val_data)
model.test()
visuals = model.get_current_visuals()
rlt_img = util.tensor2img(visuals["rlt"]) # uint8
gt_img = util.tensor2img(visuals["GT"]) # uint8
lq_img = util.tensor2img(visuals["LQ"][2]) # uint8
img_dir = opt["path"]["val_images"]
util.mkdir(img_dir)
save_img_path = os.path.join(
img_dir, "{}.png".format(idx_d))
util.save_img(np.hstack((lq_img, rlt_img, gt_img)),
save_img_path)
# calculate PSNR
psnr = util.calculate_psnr(rlt_img, gt_img)
psnr_rlt[folder].append(psnr)
pbar.update("Test {} - {}".format(folder, idx_d))
for k, v in psnr_rlt.items():
psnr_rlt_avg[k] = sum(v) / len(v)
psnr_total_avg += psnr_rlt_avg[k]
psnr_total_avg /= len(psnr_rlt)
log_s = "# Validation # PSNR: {:.4e}:".format(
psnr_total_avg)
for k, v in psnr_rlt_avg.items():
log_s += " {}: {:.4e}".format(k, v)
logger.info(log_s)
if opt["use_tb_logger"] and "debug" not in opt["name"]:
tb_logger.add_scalar("psnr_avg", psnr_total_avg,
current_step)
for k, v in psnr_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
# save models and training states
if current_step % opt["logger"]["save_checkpoint_freq"] == 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.")
tb_logger.close()
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)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# train from scratch OR resume training
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))
# config loggers. Before it, the log will not work
util.setup_logger(None,
opt['path']['log'],
'train',
level=logging.INFO,
screen=True)
util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO)
logger = logging.getLogger('base')
if resume_state:
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
option.check_resume(opt) # check resume options
logger.info(option.dict2str(opt))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
from tensorboardX import SummaryWriter
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
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
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']))
logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
total_epochs, 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)
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 = create_model(opt)
# resume training
if resume_state:
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):
for _, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
# update learning rate
model.update_learning_rate()
# training
model.feed_data(train_data)
model.optimize_parameters(current_step)
# log
if current_step % opt['logger']['print_freq'] == 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)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar(k, v, current_step)
logger.info(message)
# validation
if current_step % opt['train']['val_freq'] == 0:
avg_psnr = 0.0
avg_IS = 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 IS
IS = model.get_IS()
avg_IS += IS
# calculate PSNR
crop_size = opt['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
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.calculate_psnr(cropped_sr_img * 255,
cropped_gt_img * 255)
avg_psnr = avg_psnr / idx
avg_IS = avg_IS / idx
# log
logger.info('# Validation # PSNR: {:.4e} IS : {:.4e}'.format(
avg_psnr, avg_IS))
logger_val = logging.getLogger('val') # validation logger
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e} is： {:.4e} '.
format(epoch, current_step, avg_psnr, avg_IS))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
tb_logger.add_scalar('IS', avg_IS, current_step)
# save models and training states
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training.')
def main():
#### setup options of three networks
parser = argparse.ArgumentParser()
parser.add_argument('-opt_P',
type=str,
help='Path to option YMAL file of Predictor.')
parser.add_argument('-opt_C',
type=str,
help='Path to option YMAL file of Corrector.')
parser.add_argument('-opt_F',
type=str,
help='Path to option YMAL file of SFTMD_Net.')
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_P = option.parse(args.opt_P, is_train=True)
opt_C = option.parse(args.opt_C, is_train=True)
opt_F = option.parse(args.opt_F, is_train=True)
# convert to NoneDict, which returns None for missing keys
opt_P = option.dict_to_nonedict(opt_P)
opt_C = option.dict_to_nonedict(opt_C)
opt_F = option.dict_to_nonedict(opt_F)
# choose small opt for SFTMD test, fill path of pre-trained model_F
opt_F = opt_F['sftmd']
# create PCA matrix of enough kernel
batch_ker = util.random_batch_kernel(batch=30000,
l=opt_P['kernel_size'],
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=3,
tensor=False)
print('batch kernel shape: {}'.format(batch_ker.shape))
b = np.size(batch_ker, 0)
batch_ker = batch_ker.reshape((b, -1))
pca_matrix = util.PCA(batch_ker, k=opt_P['code_length']).float()
print('PCA matrix shape: {}'.format(pca_matrix.shape))
#### distributed training settings
if args.launcher == 'none': # disabled distributed training
opt_P['dist'] = False
opt_F['dist'] = False
opt_C['dist'] = False
rank = -1
print('Disabled distributed training.')
else:
opt_P['dist'] = True
opt_F['dist'] = True
opt_C['dist'] = True
init_dist()
world_size = torch.distributed.get_world_size(
) #Returns the number of processes in the current process group
rank = torch.distributed.get_rank(
) #Returns the rank of current process group
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
###### Predictor&Corrector train ######
#### loading resume state if exists
if opt_P['path'].get('resume_state', None):
# distributed resuming: all load into default GPU
device_id = torch.cuda.current_device()
resume_state = torch.load(
opt_P['path']['resume_state'],
map_location=lambda storage, loc: storage.cuda(device_id))
option.check_resume(opt_P,
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-7)
if resume_state is None:
# Predictor path
util.mkdir_and_rename(
opt_P['path']
['experiments_root']) # rename experiment folder if exists
util.mkdirs(
(path for key, path in opt_P['path'].items()
if not key == 'experiments_root'
and 'pretrain_model' not in key and 'resume' not in key))
# Corrector path
util.mkdir_and_rename(
opt_C['path']
['experiments_root']) # rename experiment folder if exists
util.mkdirs(
(path for key, path in opt_C['path'].items()
if not key == 'experiments_root'
and 'pretrain_model' not in key and 'resume' not in key))
# config loggers. Before it, the log will not work
util.setup_logger('base',
opt_P['path']['log'],
'train_' + opt_P['name'],
level=logging.INFO,
screen=True,
tofile=True)
util.setup_logger('val',
opt_P['path']['log'],
'val_' + opt_P['name'],
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt_P))
logger.info(option.dict2str(opt_C))
# tensorboard logger
if opt_P['use_tb_logger'] and 'debug' not in opt_P['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
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt_P['name'])
else:
util.setup_logger('base',
opt_P['path']['log'],
'train',
level=logging.INFO,
screen=True)
logger = logging.getLogger('base')
#### random seed
seed = opt_P['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.benchmark = 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_P['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_P['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
if opt_P['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_P,
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_P, 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
assert val_loader is not None
#### create model
model_F = create_model(opt_F) #load pretrained model of SFTMD
model_P = create_model(opt_P)
model_C = create_model(opt_C)
#### 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_P.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_P['dist']:
train_sampler.set_epoch(epoch)
for _, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
#### update learning rate, schedulers
# model.update_learning_rate(current_step, warmup_iter=opt_P['train']['warmup_iter'])
#### preprocessing for LR_img and kernel map
prepro = util.SRMDPreprocessing(opt_P['scale'],
pca_matrix,
para_input=opt_P['code_length'],
kernel=opt_P['kernel_size'],
noise=False,
cuda=True,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=3,
rate_cln=0.2,
noise_high=0.0)
LR_img, ker_map = prepro(train_data['GT'])
#### training Predictor
model_P.feed_data(LR_img, ker_map)
model_P.optimize_parameters(current_step)
P_visuals = model_P.get_current_visuals()
est_ker_map = P_visuals['Batch_est_ker_map']
#### log of model_P
if current_step % opt_P['logger']['print_freq'] == 0:
logs = model_P.get_current_log()
message = 'Predictor <epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step, model_P.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt_P['use_tb_logger'] and 'debug' not in opt_P['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
#### training Corrector
for step in range(opt_C['step']):
# test SFTMD for corresponding SR image
model_F.feed_data(train_data, LR_img, est_ker_map)
model_F.test()
F_visuals = model_F.get_current_visuals()
SR_img = F_visuals['Batch_SR']
# Test SFTMD to produce SR images
# train corrector given SR image and estimated kernel map
model_C.feed_data(SR_img, est_ker_map, ker_map)
model_C.optimize_parameters(current_step)
C_visuals = model_C.get_current_visuals()
est_ker_map = C_visuals['Batch_est_ker_map']
#### log of model_C
if current_step % opt_C['logger']['print_freq'] == 0:
logs = model_C.get_current_log()
message = 'Corrector <epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step,
model_C.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt_C['use_tb_logger'] and 'debug' not in opt_C[
'name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
# validation, to produce ker_map_list(fake)
if current_step % opt_P['train']['val_freq'] == 0 and rank <= 0:
avg_psnr = 0.0
idx = 0
for _, val_data in enumerate(val_loader):
prepro = util.SRMDPreprocessing(
opt_P['scale'],
pca_matrix,
para_input=opt_P['code_length'],
kernel=opt_P['kernel_size'],
noise=False,
cuda=True,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
scaling=3,
rate_cln=0.2,
noise_high=0.0)
LR_img, ker_map = prepro(val_data['GT'])
single_img_psnr = 0.0
# valid Predictor
model_P.feed_data(LR_img, ker_map)
model_P.test()
P_visuals = model_P.get_current_visuals()
est_ker_map = P_visuals['Batch_est_ker_map']
for step in range(opt_C['step']):
step += 1
idx += 1
model_F.feed_data(val_data, LR_img, est_ker_map)
model_F.test()
F_visuals = model_F.get_current_visuals()
SR_img = F_visuals['Batch_SR']
# Test SFTMD to produce SR images
model_C.feed_data(SR_img, est_ker_map, ker_map)
model_C.test()
C_visuals = model_C.get_current_visuals()
est_ker_map = C_visuals['Batch_est_ker_map']
sr_img = util.tensor2img(F_visuals['SR']) # uint8
gt_img = util.tensor2img(F_visuals['GT']) # uint8
# Save SR images for reference
img_name = os.path.splitext(
os.path.basename(val_data['LQ_path'][0]))[0]
img_dir = os.path.join(opt_P['path']['val_images'],
img_name)
# img_dir = os.path.join(opt_F['path']['val_images'], str(current_step), '_', str(step))
util.mkdir(img_dir)
save_img_path = os.path.join(
img_dir, '{:s}_{:d}_{:d}.png'.format(
img_name, current_step, step))
util.save_img(sr_img, save_img_path)
# calculate PSNR
crop_size = opt_P['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
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, :]
step_psnr = util.calculate_psnr(
cropped_sr_img * 255, cropped_gt_img * 255)
logger.info(
'<epoch:{:3d}, iter:{:8,d}, step:{:3d}> img:{:s}, psnr: {:.4f}'
.format(epoch, current_step, step, img_name,
step_psnr))
single_img_psnr += step_psnr
avg_psnr += util.calculate_psnr(
cropped_sr_img * 255, cropped_gt_img * 255)
avg_signle_img_psnr = single_img_psnr / step
logger.info(
'<epoch:{:3d}, iter:{:8,d}, step:{:3d}> img:{:s}, average psnr: {:.4f}'
.format(epoch, current_step, step, img_name,
avg_signle_img_psnr))
avg_psnr = avg_psnr / idx
# log
logger.info('# Validation # PSNR: {:.4f}'.format(avg_psnr))
logger_val = logging.getLogger('val') # validation logger
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}, step:{:3d}> psnr: {:.4f}'.
format(epoch, current_step, step, avg_psnr))
# tensorboard logger
if opt_P['use_tb_logger'] and 'debug' not in opt_P['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
#### save models and training states
if current_step % opt_P['logger']['save_checkpoint_freq'] == 0:
if rank <= 0:
logger.info('Saving models and training states.')
model_P.save(current_step)
model_P.save_training_state(epoch, current_step)
model_C.save(current_step)
model_C.save_training_state(epoch, current_step)
if rank <= 0:
logger.info('Saving the final model.')
model_P.save('latest')
model_C.save('latest')
logger.info('End of Predictor and Corrector training.')
tb_logger.close()
def SFTMD_train(opt_F, rank, world_size, pca_matrix):
#### loading resume state if exists
if opt_F['path'].get('resume_state', None):
# distributed resuming: all load into default GPU
device_id = torch.cuda.current_device()
resume_state = torch.load(opt_F['path']['resume_state'],
map_location=lambda storage, loc: storage.cuda(device_id))
option.check_resume(opt_F, resume_state['iter']) # check resume options
else:
resume_state = None
#### mkdir and loggers
if rank <= 0:
if resume_state is None:
util.mkdir_and_rename(
opt_F['path']['experiments_root']) # rename experiment folder if exists
util.mkdirs((path for key, path in opt_F['path'].items() if not key == 'experiments_root'
and 'pretrain_model' not in key and 'resume' not in key))
# config loggers. Before it, the log will not work
util.setup_logger('base', opt_F['path']['log'], 'train_' + opt_F['name'], level=logging.INFO,
screen=True, tofile=True)
util.setup_logger('val', opt_F['path']['log'], 'val_' + opt_F['name'], level=logging.INFO,
screen=True, tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt_F))
# tensorboard logger
if opt_F['use_tb_logger'] and 'debug' not in opt_F['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
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt_F['name'])
else:
util.setup_logger('base', opt_F['path']['log'], 'train', level=logging.INFO, screen=True)
logger = logging.getLogger('base')
#### create train and val dataloader
dataset_ratio = 200 # enlarge the size of each epoch
for phase, dataset_opt in opt_F['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_F['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
if opt_F['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_F, 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_F, 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
assert val_loader is not None
#### create model
model_F = create_model(opt_F)
#### 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_F.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_F['dist']:
train_sampler.set_epoch(epoch)
for _, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
#### preprocessing for LR_img and kernel map
prepro = util.SRMDPreprocessing(opt_F['scale'], pca_matrix, para_input=10, kernel=21, noise=False, cuda=True,
sig_min=0.2, sig_max=4.0, rate_iso=1.0, scaling=3,
rate_cln=0.2, noise_high=0.0)
LR_img, ker_map = prepro(train_data['GT'])
#### update learning rate, schedulers
model_F.update_learning_rate(current_step, warmup_iter=opt_F['train']['warmup_iter'])
#### training
model_F.feed_data(train_data, LR_img, ker_map)
model_F.optimize_parameters(current_step)
#### log
if current_step % opt_F['logger']['print_freq'] == 0:
logs = model_F.get_current_log()
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step, model_F.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt_F['use_tb_logger'] and 'debug' not in opt_F['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
# validation
if current_step % opt_F['train']['val_freq'] == 0 and rank <= 0:
avg_psnr = 0.0
idx = 0
for _, val_data in enumerate(val_loader):
idx += 1
#### preprocessing for LR_img and kernel map
prepro = util.SRMDPreprocessing(opt_F['scale'], pca_matrix, para_input=15, noise=False, cuda=True,
sig_min=0.2, sig_max=4.0, rate_iso=1.0, scaling=3,
rate_cln=0.2, noise_high=0.0)
LR_img, ker_map = prepro(val_data['GT'])
model_F.feed_data(val_data, LR_img, ker_map)
model_F.test()
visuals = model_F.get_current_visuals()
sr_img = util.tensor2img(visuals['SR']) # uint8
gt_img = util.tensor2img(visuals['GT']) # uint8
# Save SR images for reference
img_name = os.path.splitext(os.path.basename(val_data['LQ_path'][0]))[0]
#img_dir = os.path.join(opt_F['path']['val_images'], img_name)
img_dir = os.path.join(opt_F['path']['val_images'], str(current_step))
util.mkdir(img_dir)
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
crop_size = opt_F['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
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.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255)
avg_psnr = avg_psnr / idx
# log
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
logger_val = logging.getLogger('val') # validation logger
logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}'.format(epoch, current_step, avg_psnr))
# tensorboard logger
if opt_F['use_tb_logger'] and 'debug' not in opt_F['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
#### save models and training states
if current_step % opt_F['logger']['save_checkpoint_freq'] == 0:
if rank <= 0:
logger.info('Saving models and training states.')
model_F.save(current_step)
model_F.save_training_state(epoch, current_step)
if rank <= 0:
logger.info('Saving the final model.')
model_F.save('latest')
logger.info('End of SFTMD 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)
# Convert to NoneDict, which return None for missing key.
opt = option.dict_to_nonedict(opt)
pytorch_ver = get_pytorch_ver()
# train from scratch OR resume training
if opt['path']['resume_state']:
if os.path.isdir(opt['path']['resume_state']):
import glob
resume_state_path = util.sorted_nicely(
glob.glob(
os.path.normpath(opt['path']['resume_state']) +
'/*.state'))[-1]
else:
resume_state_path = opt['path']['resume_state']
resume_state = torch.load(resume_state_path)
else: # training from scratch
resume_state = None
# rename old folder if exists
util.mkdir_and_rename(opt['path']['experiments_root'])
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))
# config loggers. Before it, the log will not work
util.setup_logger(None,
opt['path']['log'],
'train',
level=logging.INFO,
screen=True)
util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO)
logger = logging.getLogger('base')
if resume_state:
logger.info('Set [resume_state] to ' + resume_state_path)
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
option.check_resume(opt) # check resume options
logger.info(option.dict2str(opt))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
from tensorboardX import SummaryWriter
try:
# for version tensorboardX >= 1.7
tb_logger = SummaryWriter(logdir='../tb_logger/' + opt['name'])
except:
# for version tensorboardX < 1.6
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
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
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']))
logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
total_epochs, 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)
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 = create_model(opt)
# resume training
if resume_state:
start_epoch = resume_state['epoch']
current_step = resume_state['iter']
model.resume_training(resume_state) # handle optimizers and schedulers
# updated schedulers in case JSON configuration has changed
model.update_schedulers(opt['train'])
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):
for n, train_data in enumerate(train_loader, start=1):
current_step += 1
if current_step > total_iters:
break
if pytorch_ver == "pre": # Order for PyTorch ver < 1.1.0
# update learning rate
model.update_learning_rate(current_step - 1)
# training
model.feed_data(train_data)
model.optimize_parameters(current_step)
elif pytorch_ver == "post": # Order for PyTorch ver > 1.1.0
# training
model.feed_data(train_data)
model.optimize_parameters(current_step)
# update learning rate
model.update_learning_rate(current_step - 1)
else:
print('Error identifying PyTorch version. ', torch.__version__)
break
# log
if current_step % opt['logger']['print_freq'] == 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)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar(k, v, current_step)
logger.info(message)
# save models and training states (changed to save models before validation)
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
model.save(current_step)
model.save_training_state(epoch + (n >= len(train_loader)),
current_step)
logger.info('Models and training states saved.')
# validation
if current_step % opt['train']['val_freq'] == 0:
avg_psnr = 0.0
avg_ssim = 0.0
avg_lpips = 0.0
idx = 0
val_sr_imgs_list = []
val_gt_imgs_list = []
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()
if opt['datasets']['train'][
'znorm']: # If the image range is [-1,1]
sr_img = util.tensor2img(visuals['SR'],
min_max=(-1, 1)) # uint8
gt_img = util.tensor2img(visuals['HR'],
min_max=(-1, 1)) # uint8
else: # Default: Image range is [0,1]
sr_img = util.tensor2img(visuals['SR']) # uint8
gt_img = util.tensor2img(visuals['HR']) # uint8
# sr_img = util.tensor2img(visuals['SR']) # uint8
# gt_img = util.tensor2img(visuals['HR']) # uint8
# print("Min. SR value:",sr_img.min()) # Debug
# print("Max. SR value:",sr_img.max()) # Debug
# print("Min. GT value:",gt_img.min()) # Debug
# print("Max. GT value:",gt_img.max()) # Debug
# 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, SSIM and LPIPS distance
crop_size = opt['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
# For training models with only one channel ndim==2, if RGB ndim==3, etc.
if gt_img.ndim == 2:
cropped_gt_img = gt_img[crop_size:-crop_size,
crop_size:-crop_size]
else:
cropped_gt_img = gt_img[crop_size:-crop_size,
crop_size:-crop_size, :]
if sr_img.ndim == 2:
cropped_sr_img = sr_img[crop_size:-crop_size,
crop_size:-crop_size]
else: # Default: RGB images
cropped_sr_img = sr_img[crop_size:-crop_size,
crop_size:-crop_size, :]
# If calculating only once for all images
val_gt_imgs_list.append(cropped_gt_img)
# If calculating only once for all images
val_sr_imgs_list.append(cropped_sr_img)
# LPIPS only works for RGB images
avg_psnr += util.calculate_psnr(cropped_sr_img * 255,
cropped_gt_img * 255)
avg_ssim += util.calculate_ssim(cropped_sr_img * 255,
cropped_gt_img * 255)
# If calculating for each image
# avg_lpips += lpips.calculate_lpips([cropped_sr_img], [cropped_gt_img])
avg_psnr = avg_psnr / idx
avg_ssim = avg_ssim / idx
# avg_lpips=avg_lpips / idx # If calculating for each image
# If calculating only once for all images
avg_lpips = lpips.calculate_lpips(val_sr_imgs_list,
val_gt_imgs_list)
# log
# logger.info('# Validation # PSNR: {:.5g}, SSIM: {:.5g}'.format(avg_psnr, avg_ssim))
logger.info(
'# Validation # PSNR: {:.5g}, SSIM: {:.5g}, LPIPS: {:.5g}'.
format(avg_psnr, avg_ssim, avg_lpips))
logger_val = logging.getLogger('val') # validation logger
# logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.5g}, ssim: {:.5g}'.format(
# epoch, current_step, avg_psnr, avg_ssim))
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr: {:.5g}, ssim: {:.5g}, lpips: {:.5g}'
.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('psnr', avg_psnr, current_step)
tb_logger.add_scalar('ssim', avg_ssim, current_step)
tb_logger.add_scalar('lpips', avg_lpips, current_step)
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training.')
def main():
#### options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
default='options/train/train_EDVR_woTSA_M.yml',
help='Path to option YAML file.')
parser.add_argument('--set',
dest='set_opt',
default=None,
nargs=argparse.REMAINDER,
help='set options')
args = parser.parse_args()
opt = option.parse(args.opt, args.set_opt, is_train=True)
#### loading resume state if exists
if opt['path'].get('resume_state', None):
# distributed resuming: all load into default GPU
print('Training from state: {}'.format(opt['path']['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))
option.check_resume(opt, resume_state['iter']) # check resume options
elif opt['auto_resume']:
exp_dir = opt['path']['experiments_root']
# first time run: create dirs
if not os.path.exists(exp_dir):
os.makedirs(exp_dir)
os.makedirs(opt['path']['models'])
os.makedirs(opt['path']['training_state'])
os.makedirs(opt['path']['val_images'])
os.makedirs(opt['path']['tb_logger'])
resume_state = None
else:
# detect experiment directory and get the latest state
state_dir = opt['path']['training_state']
state_files = [
x for x in os.listdir(state_dir) if x.endswith('state')
]
# no valid state detected
if len(state_files) < 1:
print(
'No previous training state found, train from start state')
resume_state = None
else:
state_files = sorted(state_files,
key=lambda x: int(x.split('.')[0]))
latest_state = state_files[-1]
print('Training from lastest state: {}'.format(latest_state))
latest_state_file = os.path.join(state_dir, latest_state)
opt['path']['resume_state'] = latest_state_file
device_id = torch.cuda.current_device()
resume_state = torch.load(
latest_state_file,
map_location=lambda storage, loc: storage.cuda(device_id))
option.check_resume(opt, resume_state['iter'])
else:
resume_state = None
if resume_state is None and not opt['auto_resume'] and not opt['no_log']:
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))
# 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)
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.2: # 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
tb_logger = SummaryWriter(log_dir=opt['path']['tb_logger'])
# convert to NoneDict, which returns None for missing keys
opt = option.dict_to_nonedict(opt)
#### random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
logger.info('Random seed: {}'.format(seed))
util.set_random_seed(seed)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
#### create train and val dataloader
if opt['datasets']['train']['ratio']:
dataset_ratio = opt['datasets']['train']['ratio']
else:
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 * dataset_ratio)))
if dataset_opt['mode'] in ['MetaREDS', 'MetaREDSOnline']:
train_sampler = MetaIterSampler(train_set,
dataset_opt['batch_size'],
len(opt['scale']),
dataset_ratio)
elif dataset_opt['mode'] in ['REDS', 'MultiREDS']:
train_sampler = IterSampler(train_set,
dataset_opt['batch_size'],
dataset_ratio)
else:
train_sampler = None
train_loader = create_dataloader(train_set, dataset_opt, opt,
train_sampler)
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)
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))
#### create model
model = create_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):
train_sampler.set_epoch(epoch)
for _, train_data in enumerate(train_loader):
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)
#### log
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
message = '[epoch:{:3d}, iter:{:8,d}, lr:('.format(
epoch, current_step)
for v in model.get_current_learning_rate():
message += '{:.3e},'.format(v)
message += ')] '
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar(k, v, current_step)
logger.info(message)
print("PROGRESS: {:02d}%".format(
int(current_step / total_iters * 100)))
#### validation
if opt['datasets'].get(
'val',
None) and current_step % opt['train']['val_freq'] == 0:
pbar = util.ProgressBar(len(val_loader))
psnr_rlt = {} # with border and center frames
psnr_rlt_avg = {}
psnr_total_avg = 0.
for val_data in val_loader:
folder = val_data['folder'][0]
idx_d = val_data['idx'].item()
# border = val_data['border'].item()
if psnr_rlt.get(folder, None) is None:
psnr_rlt[folder] = []
model.feed_data(val_data)
model.test()
visuals = model.get_current_visuals()
rlt_img = util.tensor2img(visuals['rlt']) # uint8
gt_img = util.tensor2img(visuals['GT']) # uint8
# calculate PSNR
psnr = util.calculate_psnr(rlt_img, gt_img)
psnr_rlt[folder].append(psnr)
pbar.update('Test {} - {}'.format(folder, idx_d))
for k, v in psnr_rlt.items():
psnr_rlt_avg[k] = sum(v) / len(v)
psnr_total_avg += psnr_rlt_avg[k]
psnr_total_avg /= len(psnr_rlt)
log_s = '# Validation # PSNR: {:.4e}:'.format(psnr_total_avg)
for k, v in psnr_rlt_avg.items():
log_s += ' {}: {:.4e}'.format(k, v)
logger.info(log_s)
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr_avg', psnr_total_avg,
current_step)
for k, v in psnr_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
#### save models and training states
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training.')
tb_logger.close()
def SFTMD_test(opt_F):
# config loggers. Before it, the log will not work
util.setup_logger('base',
opt_F['path']['log'],
'test_' + opt_F['name'],
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt_F))
for phase, dataset_opt in sorted(opt_F['datasets'].items()):
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set, dataset_opt, opt_F)
logger.info('Number of test images in [{:s}]: {:d}'.format(
dataset_opt['name'], len(test_set)))
model_F = create_model(opt_F)
logger.info('\nTesting [{:s}]...'.format(opt_F['model']))
dataset_dir = os.path.join(opt_F['path']['results_root'],
dataset_opt['name'])
util.mkdir(dataset_dir)
avg_psnr = 0.0
idx = 0
for _, data in enumerate(test_loader):
idx += 1
model_F.feed_data(data)
model_F.test()
visuals = model_F.get_current_visuals()
gt_img = util.tensor2img(visuals['GT']) # uint8
sr_img = util.tensor2img(visuals['SR']) # uint8
# Save SR images for reference
img_name = os.path.splitext(os.path.basename(
data['LQ_path'][0]))[0]
img_dir = os.path.join(dataset_dir, img_name)
util.mkdir(img_dir)
save_img_path = os.path.join(img_dir, '{}.png'.format(img_name))
util.save_img(sr_img, save_img_path)
# calculate PSNR
crop_size = opt_F['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
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.calculate_psnr(cropped_sr_img * 255,
cropped_gt_img * 255)
avg_psnr = avg_psnr / idx
# log
logger.info('# test # PSNR: {:.4e}'.format(avg_psnr))
print('SR image of SFTMD has been saved to list')
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)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# train from scratch OR resume training
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))
# config loggers. Before it, the log will not work
util.setup_logger(None,
opt['path']['log'],
'train',
level=logging.INFO,
screen=True)
util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO)
logger = logging.getLogger('base')
if resume_state:
# resume_state[]
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
option.check_resume(opt) # check resume options
logger.info(option.dict2str(opt))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
from tensorboardX import SummaryWriter
tb_logger = SummaryWriter(logdir='../../SRN_tb_logger/' + opt['name'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
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
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']))
logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
total_epochs, 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)
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 = create_model(opt)
# resume training
if resume_state:
start_epoch = resume_state['epoch']
current_step = resume_state['iter']
model.resume_training(resume_state,
opt['train']) # 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):
for _, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
# update learning rate
model.update_learning_rate()
# training
model.feed_data(train_data, True)
model.optimize_parameters(current_step)
# log
if current_step % opt['logger']['print_freq'] == 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)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar(k, v, current_step)
logger.info(message)
# training samples
if opt['train']['save_tsamples'] and current_step % opt['train'][
'save_tsamples'] == 0:
fake_LRs = os.listdir(
opt['datasets']['train']['dataroot_fake_LR'])
real_LRs = os.listdir(
opt['datasets']['train']['dataroot_real_LR'])
HRs = os.listdir(opt['datasets']['train']['dataroot_HR'])
for i in range(5):
random_index = np.random.choice(range(len(fake_LRs)))
fake_LR_path = os.path.join(
opt['datasets']['train']['dataroot_fake_LR'],
fake_LRs[random_index])
real_LR_path = os.path.join(
opt['datasets']['train']['dataroot_real_LR'],
real_LRs[random_index])
HR_path = os.path.join(
opt['datasets']['train']['dataroot_HR'],
HRs[random_index])
fake_LR = np.array(Image.open(fake_LR_path))
real_LR = np.array(Image.open(real_LR_path))
HR = np.array(Image.open(HR_path))
h, w, _ = fake_LR.shape
fake_LR = fake_LR[h // 2 - 64:h // 2 + 64,
w // 2 - 64:w // 2 + 64, :]
h, w, _ = HR.shape
HR = HR[h // 2 - 64 * 4:h // 2 + 64 * 4,
w // 2 - 64 * 4:w // 2 + 64 * 4, :]
h, w, _ = real_LR.shape
real_LR = real_LR[h // 2 - 64:h // 2 + 64,
w // 2 - 64:w // 2 + 64, :]
fake_LR = torch.from_numpy(
np.ascontiguousarray(np.transpose(
fake_LR, (2, 0, 1)))).float().unsqueeze(0) / 255
real_LR = torch.from_numpy(
np.ascontiguousarray(np.transpose(
real_LR, (2, 0, 1)))).float().unsqueeze(0) / 255
HR = torch.from_numpy(
np.ascontiguousarray(np.transpose(
HR, (2, 0, 1)))).float().unsqueeze(0) / 255
LR = torch.cat([fake_LR, real_LR], dim=0)
data = {'LR': LR, 'HR': HR}
model.feed_data(data, False)
model.test(tsamples=True)
visuals = model.get_current_visuals(tsamples=True)
fake_SR = visuals['SR'][0]
real_SR = visuals['SR'][1]
fake_hf = visuals['hf'][0]
real_hf = visuals['hf'][1]
HR = visuals['HR']
HR_hf = visuals['HR_hf'][0]
# image_1 = torch.cat([fake_LR[0], fake_SR[0]], dim=2)
# image_2 = torch.cat([real_LR[0], real_SR[0]], dim=2)
image_1 = np.clip(torch.cat([fake_SR, HR, real_SR], dim=2),
0, 1)
image_2 = np.clip(
torch.cat([fake_hf, HR_hf, real_hf], dim=2), 0, 1)
image = torch.cat([image_1, image_2], dim=1)
tb_logger.add_image(
'train/train_samples_{}'.format(str(i)), image,
current_step)
logger.info('Saved training Samples')
# validation
if current_step % opt['train']['val_freq'] == 0:
avg_psnr = 0.0
idx = 0
avg_lpips = 0.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, False)
model.test()
visuals = model.get_current_visuals()
sr_img = util.tensor2img(visuals['SR']) # uint8
if 'HR' in opt['datasets']['val']['mode']:
gt_img = util.tensor2img(visuals['HR']) # uint8
log_info = '{}'.format(
val_data['HR_path'][0].split('/')[-1])
if opt['val_lpips']:
lpips = visuals['LPIPS']
avg_lpips += lpips
log_info += ' LPIPS:{:.3f}'.format(
lpips.numpy())
if opt['use_domain_distance_map']:
ada_w = visuals['adaptive_weights']
log_info += ' Adaptive weights:{:.2f}'.format(
ada_w.numpy())
# logger.info('{} LPIPS: {:.3f}'.format(val_data['HR_path'][0].split('/')[-1], lpips.numpy()))
# print('img:', val_data['HR_path'][0].split('/')[-1], 'LPIPS: %.3f' % lpips.numpy())
# else:
# print('img:', val_data['LR_path'][0].split('/')[-1])
logger.info(log_info)
# 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 'HR' in opt['datasets']['val']['mode']:
crop_size = opt['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
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.calculate_psnr(
cropped_sr_img * 255, cropped_gt_img * 255)
avg_psnr = avg_psnr / idx
if opt['val_lpips']:
avg_lpips = avg_lpips / idx
print('Mean LPIPS:', avg_lpips.numpy())
# log
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
logger_val = logging.getLogger('val') # validation logger
if opt['val_lpips']:
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}, LPIPS: {:.4f}'
.format(epoch, current_step, avg_psnr, avg_lpips))
else:
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}'.format(
epoch, current_step, avg_psnr))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
tb_logger.add_scalar('LPIPS', avg_lpips, current_step)
# save models and training states
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training.')
def main():
############################################
#
# set options
#
############################################
parser = argparse.ArgumentParser()
parser.add_argument('--opt', type=str, help='Path to option YAML 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()
print("Rank:", rank)
print("World Size", world_size)
print("------------------DIST-------------------------")
############################################
#
# 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 'debug' in opt['name']:
debug_mode = True
else:
debug_mode = False
if rank <= 0: # normal training (rank -1) OR distributed training (rank 0)
if resume_state is None:
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))
# 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('base_val',
opt['path']['log'],
'val_' + opt['name'],
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
logger_val = logging.getLogger('base_val')
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
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
else:
# config loggers. Before it, the log will not work
util.setup_logger('base',
opt['path']['log'],
'train_',
level=logging.INFO,
screen=True)
print("set train log")
util.setup_logger('base_val',
opt['path']['log'],
'val_',
level=logging.INFO,
screen=True)
print("set val log")
logger = logging.getLogger('base')
logger_val = logging.getLogger('base_val')
# convert to NoneDict, which returns None for missing keys
opt = option.dict_to_nonedict(opt)
#### 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.benchmark = True
# torch.backends.cudnn.deterministic = True
############################################
#
# create train and val dataloader
#
############################################
####
# dataset_ratio = 200 # enlarge the size of each epoch
dataset_ratio = 200 # enlarge the size of each epoch
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
if opt['datasets']['train'].get('split', None):
train_set, val_set = create_dataset(dataset_opt)
else:
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))
total_iters = train_size
total_epochs = int(opt['train']['epoch'])
if opt['dist']:
train_sampler = DistIterSampler(train_set, world_size, rank,
dataset_ratio)
# total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio)))
total_epochs = int(opt['train']['epoch'])
if opt['train']['enable'] == False:
total_epochs = 1
else:
# train_sampler = None
train_sampler = RandomBalancedSampler(train_set, train_size)
train_loader = create_dataloader(train_set,
dataset_opt,
opt,
train_sampler,
vscode_debug=debug_mode)
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':
if not opt['datasets']['train'].get('split', None):
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set,
dataset_opt,
opt,
None,
vscode_debug=debug_mode)
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 = create_model(opt)
print("Model Created! ")
#### 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
print("Not Resume Training")
############################################
#
# training
#
############################################
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(
start_epoch, current_step))
model.train_AverageMeter()
saved_total_loss = 10e10
saved_total_PSNR = -1
saved_total_SSIM = -1
for epoch in range(start_epoch, total_epochs):
############################################
#
# Start a new epoch
#
############################################
current_step = 0
if opt['dist']:
train_sampler.set_epoch(epoch)
for train_idx, train_data in enumerate(train_loader):
# print('current_step', current_step)
if 'debug' in opt['name']:
img_dir = os.path.join(opt['path']['train_images'])
util.mkdir(img_dir)
LQs = train_data['LQs'] # B N C H W
if not 'sr' in opt['name']:
GTenh = train_data['GTenh']
GTinp = train_data['GTinp']
for imgs, name in zip([LQs, GTenh, GTinp],
['LQs', 'GTenh', 'GTinp']):
num = imgs.size(1)
for i in range(num):
img = util.tensor2img(imgs[0, i, ...]) # uint8
save_img_path = os.path.join(
img_dir, '{:4d}_{:s}_{:1d}.png'.format(
train_idx, str(name), i))
util.save_img(img, save_img_path)
else:
if 'GT' in train_data:
GT_name = 'GT'
elif 'GTs' in train_data:
GT_name = 'GTs'
GT = train_data[GT_name]
for imgs, name in zip([LQs, GT], ['LQs', GT_name]):
if name == 'GT':
num = imgs.size(0)
img = util.tensor2img(imgs[0, ...]) # uint8
save_img_path = os.path.join(
img_dir, '{:4d}_{:s}_{:1d}.png'.format(
train_idx, str(name), 0))
util.save_img(img, save_img_path)
elif name == 'GTs':
num = imgs.size(1)
for i in range(num):
img = util.tensor2img(imgs[:, i, ...]) # uint8
save_img_path = os.path.join(
img_dir, '{:4d}_{:s}_{:1d}.png'.format(
train_idx, str(name), i))
util.save_img(img, save_img_path)
else:
num = imgs.size(1)
for i in range(num):
img = util.tensor2img(imgs[:, i, ...]) # uint8
save_img_path = os.path.join(
img_dir, '{:4d}_{:s}_{:1d}.png'.format(
train_idx, str(name), i))
util.save_img(img, save_img_path)
if (train_idx >= 3): # set to 0, just do validation
break
# if pre-load weight first do validation and skip the first epoch
# if opt['path'].get('pretrain_model_G', None) and epoch == 0:
# epoch += 1
# break
if opt['train']['enable'] == False:
message_train_loss = 'None'
break
current_step += 1
if current_step > total_iters:
print("Total Iteration Reached !")
break
#### update learning rate
if opt['train']['lr_scheme'] == 'ReduceLROnPlateau':
pass
else:
model.update_learning_rate(
current_step, warmup_iter=opt['train']['warmup_iter'])
#### training
model.feed_data(train_data)
model.optimize_parameters(current_step)
model.train_AverageMeter_update()
#### log
if current_step % opt['logger']['print_freq'] == 0:
logs_inst, logs_avg = model.get_current_log(
) # training loss mode='train'
message = '[epoch:{:3d}, iter:{:8,d}, lr:('.format(
epoch, current_step)
for v in model.get_current_learning_rate():
message += '{:.3e},'.format(v)
message += ')] '
# if 'debug' in opt['name']: # debug model print the instant loss
# for k, v in logs_inst.items():
# message += '{:s}: {:.4e} '.format(k, v)
# # tensorboard logger
# if opt['use_tb_logger'] and 'debug' not in opt['name']:
# if rank <= 0:
# tb_logger.add_scalar(k, v, current_step)
# for avg loss
current_iters_epoch = epoch * total_iters + current_step
for k, v in logs_avg.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_iters_epoch)
if rank <= 0:
logger.info(message)
# saving models
if epoch == 1:
save_filename = '{:04d}_{}.pth'.format(0, 'G')
save_path = os.path.join(opt['path']['models'], save_filename)
if os.path.exists(save_path):
os.remove(save_path)
save_filename = '{:04d}_{}.pth'.format(epoch - 1, 'G')
save_path = os.path.join(opt['path']['models'], save_filename)
if os.path.exists(save_path):
os.remove(save_path)
if rank <= 0:
logger.info('Saving models and training states.')
save_filename = '{:04d}'.format(epoch)
model.save(save_filename)
# ======================================================================= #
# Main validation loop #
# ======================================================================= #
if opt['datasets'].get('val', None):
if opt['dist']:
# multi-GPU testing
psnr_rlt = {} # with border and center frames
psnr_rlt_avg = {}
psnr_total_avg = 0.
ssim_rlt = {} # with border and center frames
ssim_rlt_avg = {}
ssim_total_avg = 0.
val_loss_rlt = {} # the averaged loss
val_loss_rlt_avg = {}
val_loss_total_avg = 0.
if rank == 0:
pbar = util.ProgressBar(len(val_set))
for idx in range(
rank, len(val_set),
world_size): # distributed parallel validation
# print('idx', idx)
if 'debug' in opt['name']:
if (idx >= 3):
break
if (idx >= 1000):
break
val_data = val_set[idx]
# use idx method to fetch must extend batch dimension
val_data['LQs'].unsqueeze_(0)
val_data['GTenh'].unsqueeze_(0)
val_data['GTinp'].unsqueeze_(0)
key = val_data['key'][0] # IMG_0034_00809
max_idx = len(val_set)
val_name = 'val_set'
num = model.get_info(
) # each model has different number of loss
if psnr_rlt.get(val_name, None) is None:
psnr_rlt[val_name] = torch.zeros([num, max_idx],
dtype=torch.float32,
device='cuda')
if ssim_rlt.get(val_name, None) is None:
ssim_rlt[val_name] = torch.zeros([num, max_idx],
dtype=torch.float32,
device='cuda')
if val_loss_rlt.get(val_name, None) is None:
val_loss_rlt[val_name] = torch.zeros(
[num, max_idx], dtype=torch.float32, device='cuda')
model.feed_data(val_data)
model.test()
avg_loss, loss_list = model.get_loss(ret=1)
save_enable = True
if idx >= 100:
save_enable = False
psnr_list, ssim_list = model.compute_current_psnr_ssim(
save=save_enable,
name=key,
save_path=opt['path']['val_images'])
# print('psnr_list',psnr_list)
assert len(loss_list) == num
assert len(psnr_list) == num
for i in range(num):
psnr_rlt[val_name][i, idx] = psnr_list[i]
ssim_rlt[val_name][i, idx] = ssim_list[i]
val_loss_rlt[val_name][i, idx] = loss_list[i]
# print('psnr_rlt[val_name][i, idx]',psnr_rlt[val_name][i, idx])
# print('ssim_rlt[val_name][i, idx]',ssim_rlt[val_name][i, idx])
# print('val_loss_rlt[val_name][i, idx] ',val_loss_rlt[val_name][i, idx] )
if rank == 0:
for _ in range(world_size):
pbar.update('Test {} - {}/{}'.format(
key, idx, max_idx))
# # collect data
for _, v in psnr_rlt.items():
for i in v:
dist.reduce(i, 0)
for _, v in ssim_rlt.items():
for i in v:
dist.reduce(i, 0)
for _, v in val_loss_rlt.items():
for i in v:
dist.reduce(i, 0)
dist.barrier()
if rank == 0:
psnr_rlt_avg = {}
psnr_total_avg = 0.
for k, v in psnr_rlt.items(): # key, value
# print('k', k, 'v', v, 'v.shape', v.shape)
psnr_rlt_avg[k] = []
for i in range(num):
non_zero_idx = v[i, :].nonzero()
# logger.info('non_zero_idx {}'.format(non_zero_idx.shape)) # check
matrix = v[i, :][non_zero_idx]
# print('matrix', matrix)
value = torch.mean(matrix).cpu().item()
# print('value', value)
psnr_rlt_avg[k].append(value)
psnr_total_avg += psnr_rlt_avg[k][i]
psnr_total_avg = psnr_total_avg / (len(psnr_rlt) * num)
log_p = '# Validation # Avg. PSNR: {:.2f},'.format(
psnr_total_avg)
for k, v in psnr_rlt_avg.items():
for i, it in enumerate(v):
log_p += ' {}: {:.2f}'.format(i, it)
logger.info(log_p)
logger_val.info(log_p)
# ssim
ssim_rlt_avg = {}
ssim_total_avg = 0.
for k, v in ssim_rlt.items():
ssim_rlt_avg[k] = []
for i in range(num):
non_zero_idx = v[i, :].nonzero()
# print('non_zero_idx', non_zero_idx)
matrix = v[i, :][non_zero_idx]
# print('matrix', matrix)
value = torch.mean(matrix).cpu().item()
# print('value', value)
ssim_rlt_avg[k].append(
torch.mean(matrix).cpu().item())
ssim_total_avg += ssim_rlt_avg[k][i]
ssim_total_avg /= (len(ssim_rlt) * num)
log_s = '# Validation # Avg. SSIM: {:.2f},'.format(
ssim_total_avg)
for k, v in ssim_rlt_avg.items():
for i, it in enumerate(v):
log_s += ' {}: {:.2f}'.format(i, it)
logger.info(log_s)
logger_val.info(log_s)
# added
val_loss_rlt_avg = {}
val_loss_total_avg = 0.
for k, v in val_loss_rlt.items():
# k, key, the folder name
# v, value, the torch matrix
val_loss_rlt_avg[k] = [] # loss0 - loss_N
for i in range(num):
non_zero_idx = v[i, :].nonzero()
# print('non_zero_idx', non_zero_idx)
matrix = v[i, :][non_zero_idx]
# print('matrix', matrix)
value = torch.mean(matrix).cpu().item()
# print('value', value)
val_loss_rlt_avg[k].append(
torch.mean(matrix).cpu().item())
val_loss_total_avg += val_loss_rlt_avg[k][i]
val_loss_total_avg /= (len(val_loss_rlt) * num)
log_l = '# Validation # Avg. Loss: {:.4e},'.format(
val_loss_total_avg)
for k, v in val_loss_rlt_avg.items():
for i, it in enumerate(v):
log_l += ' {}: {:.4e}'.format(i, it)
logger.info(log_l)
logger_val.info(log_l)
message = ''
for v in model.get_current_learning_rate():
message += '{:.5e}'.format(v)
logger_val.info(
'Epoch {:02d}, LR {:s}, PSNR {:.4f}, SSIM {:.4f}, Val Loss {:.4e}'
.format(epoch, message, psnr_total_avg, ssim_total_avg,
val_loss_total_avg))
else:
pbar = util.ProgressBar(len(val_loader))
model.val_loss_AverageMeter()
model.val_AverageMeter_para()
for val_inx, val_data in enumerate(val_loader):
# if 'debug' in opt['name']:
# if (val_inx >= 10):
# break
save_enable = True
if val_inx >= 100:
save_enable = False
if val_inx >= 100:
break
key = val_data['key'][0]
folder = key[:-6]
model.feed_data(val_data)
model.test()
avg_loss, loss_list = model.get_loss(ret=1)
model.val_loss_AverageMeter_update(loss_list, avg_loss)
psnr_list, ssim_list = model.compute_current_psnr_ssim(
save=save_enable,
name=key,
save_path=opt['path']['val_images'])
model.val_AverageMeter_para_update(psnr_list, ssim_list)
if 'debug' in opt['name']:
msg_psnr = ''
msg_ssim = ''
for i, psnr in enumerate(psnr_list):
msg_psnr += '{} :{:.02f} '.format(i, psnr)
for i, ssim in enumerate(ssim_list):
msg_ssim += '{} :{:.02f} '.format(i, ssim)
logger.info('{}_{:02d} {}'.format(
key, val_inx, msg_psnr))
logger.info('{}_{:02d} {}'.format(
key, val_inx, msg_ssim))
pbar.update('Test {} - {}'.format(key, val_inx))
# toal validation log
lr = ''
for v in model.get_current_learning_rate():
lr += '{:.5e}'.format(v)
logs_avg, logs_psnr_avg, psnr_total_avg, ssim_total_avg, val_loss_total_avg = model.get_current_log(
mode='val')
msg_logs_avg = ''
for k, v in logs_avg.items():
msg_logs_avg += '{:s}: {:.4e} '.format(k, v)
logger_val.info('Val-Epoch {:02d}, LR {:s}, {:s}'.format(
epoch, lr, msg_logs_avg))
logger.info('Val-Epoch {:02d}, LR {:s}, {:s}'.format(
epoch, lr, msg_logs_avg))
msg_logs_psnr_avg = ''
for k, v in logs_psnr_avg.items():
msg_logs_psnr_avg += '{:s}: {:.4e} '.format(k, v)
logger_val.info('Val-Epoch {:02d}, LR {:s}, {:s}'.format(
epoch, lr, msg_logs_psnr_avg))
logger.info('Val-Epoch {:02d}, LR {:s}, {:s}'.format(
epoch, lr, msg_logs_psnr_avg))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('val_psnr', psnr_total_avg, epoch)
tb_logger.add_scalar('val_loss', val_loss_total_avg, epoch)
############################################
#
# end of validation, save model
#
############################################
#
if rank <= 0:
logger.info("Finished an epoch, Check and Save the model weights")
# we check the validation loss instead of training loss. OK~
if saved_total_loss >= val_loss_total_avg:
saved_total_loss = val_loss_total_avg
#torch.save(model.state_dict(), args.save_path + "/best" + ".pth")
model.save('best')
logger.info(
"Best Weights updated for decreased validation loss")
else:
logger.info(
"Weights Not updated for undecreased validation loss")
if saved_total_PSNR <= psnr_total_avg:
saved_total_PSNR = psnr_total_avg
model.save('bestPSNR')
logger.info(
"Best Weights updated for increased validation PSNR")
else:
logger.info(
"Weights Not updated for unincreased validation PSNR")
############################################
#
# end of one epoch, schedule LR
#
############################################
model.train_AverageMeter_reset()
# add scheduler todo
if opt['train']['lr_scheme'] == 'ReduceLROnPlateau':
for scheduler in model.schedulers:
# scheduler.step(val_loss_total_avg)
scheduler.step(val_loss_total_avg)
if rank <= 0:
logger.info('Saving the final model.')
model.save('last')
logger.info('End of training.')
tb_logger.close()
def main():
#### options
parser = argparse.ArgumentParser()
parser.add_argument("-opt", type=str, help="Path to option YAML 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:
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
)
)
# 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,
)
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
tb_logger = SummaryWriter(log_dir="../tb_logger/" + opt["name"])
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)
#### 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.benchmark = 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":
pass
# val_set = create_dataset(dataset_opt, isVal=True)
# 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_path = opt["path"]["pretrain_model_G"]
model = create_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)
for _, train_data in enumerate(train_loader):
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)
#### log
if current_step % opt["logger"]["print_freq"] == 0:
logs = model.get_current_log()
message = "<epoch:{:3d}, iter:{:8,d}, lr:(".format(epoch, current_step)
for v in model.get_current_learning_rate():
message += "{:.3e},".format(v)
message += ")>"
for k, v in logs.items():
message += "{:s}: {:.4e} ".format(k, v)
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
#### validation
# currently, it does not support validation during training
# if current_step % opt["train"]["val_freq"] == 0:
# avg_psnr = 0
# idx = 0
# for val_data in val_loader:
# idx += 1
# key = (
# val_data["key"][0]
# if type(val_data["key"]) is list
# else val_data["key"]
# )
# imgName = key + ".png"
# savePath = os.path.join(
# opt["path"]["val_images"], str(current_step), imgName
# )
# model.feed_data(val_data)
# model.test()
# output = model.get_current_visuals()
# hr = util.tensor2img(output["GT"])
# sr = util.tensor2img(output["restore"])
# # Cropping to calculate PSNR
# hr /= 255.0
# sr /= 255.0
# scale = 4
# H, W, C = hr.shape
# H_r, W_r = H % scale, W % scale
# cropped_hr = hr[: H - H_r, : W - W_r, :]
# cropped_sr = sr[: H - H_r, : W - W_r, :]
# avg_psnr += util.calculate_psnr(cropped_sr * 255, cropped_hr * 255)
# logger.info("Saving output in {}".format(savePath))
# util.mkdir(savePath)
# util.save_img(
# output, joinPath(savePath, str(current_step) + ".png")
# )
# avg_psnr /= idx
# # log
# logger.info("# Validation # PSNR: {:.4e}".format(avg_psnr))
# logger_val = logging.getLogger("val") # validation logger
# logger_val.info(
# "<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}".format(
# epoch, current_step, avg_psnr
# )
# )
# # tensorboard logger
# if opt["use_tb_logger"] and "debug" not in opt["name"]:
# tb_logger.add_scalar("psnr", avg_psnr, current_step)
#### save models and training states
if current_step % opt["logger"]["save_checkpoint_freq"] == 0:
if rank <= 0:
# Save the experiments in case of Colab is timeout
logger.info("Saving models and training states.")
model.save(current_step)
model.save_training_state(epoch, current_step)
copy_tree(
"/content/EDVR/experiments",
"/content/drive/My Drive/LVTN/SuperResolution/EDVR/experiments",
)
copy_tree(
"/content/EDVR/tb_logger",
"/content/drive/My Drive/LVTN/SuperResolution/EDVR/tb_logger",
)
if rank <= 0:
logger.info("Saving the final model.")
model.save("latest")
logger.info("End of training.")
tb_logger.close()
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)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# train from scratch OR resume training
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))
# config loggers. Before it, the log will not work
util.setup_logger(None,
opt["path"]["log"],
"train",
level=logging.INFO,
screen=True)
util.setup_logger("val", opt["path"]["log"], "val", level=logging.INFO)
logger = logging.getLogger("base")
if resume_state:
logger.info("Resuming training from epoch: {}, iter: {}.".format(
resume_state["epoch"], resume_state["iter"]))
option.check_resume(opt) # check resume options
logger.info(option.dict2str(opt))
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
from tensorboardX import SummaryWriter
tb_logger = SummaryWriter(log_dir="../tb_logger/" + opt["name"])
# random seed
seed = opt["train"]["manual_seed"]
if seed is None:
seed = random.randint(1, 10000)
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
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"]))
logger.info("Number of train images: {:,d}, iters: {:,d}".format(
len(train_set), train_size))
total_iters = int(opt["train"]["niter"])
total_epochs = int(math.ceil(total_iters / train_size))
logger.info("Total epochs needed: {:d} for iters {:,d}".format(
total_epochs, 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)
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 = create_model(opt)
# resume training
if resume_state:
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):
for _, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
# update learning rate
model.update_learning_rate()
# training
model.feed_data(train_data)
model.optimize_parameters(current_step)
# log
if current_step % opt["logger"]["print_freq"] == 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)
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
tb_logger.add_scalar(k, v, current_step)
logger.info(message)
# validation
if current_step % opt["train"]["val_freq"] == 0:
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
crop_size = opt["scale"]
gt_img = gt_img / 255.0
sr_img = sr_img / 255.0
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.calculate_psnr(cropped_sr_img * 255,
cropped_gt_img * 255)
avg_psnr = avg_psnr / idx
# log
logger.info("# Validation # PSNR: {:.4e}".format(avg_psnr))
logger_val = logging.getLogger("val") # validation logger
logger_val.info(
"<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}".format(
epoch, current_step, avg_psnr))
# tensorboard logger
if opt["use_tb_logger"] and "debug" not in opt["name"]:
tb_logger.add_scalar("psnr", avg_psnr, current_step)
# save models and training states
if current_step % opt["logger"]["save_checkpoint_freq"] == 0:
logger.info("Saving models and training states.")
model.save(current_step)
model.save_training_state(epoch, current_step)
copy_tree(
opt["path"]["experiments_root"],
"/content/gdrive/My Drive/LVTN/SuperResolution/SR_models/"
+ "-ESRGAN/experiments/" + opt["name"],
)
logger.info("Saving the final model.")
model.save("latest")
logger.info("End of training.")
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, default='options/test/test_ppon.json', help='Path to options JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=False)
util.mkdirs((path for key, path in opt['path'].items() if not key == 'pretrain_model_G'))
opt = option.dict_to_nonedict(opt)
util.setup_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True)
logger = logging.getLogger('base')
logger.info(option.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)
logger.info('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']
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
model.feed_data(data, need_HR=need_HR)
img_path = data['LR_path'][0]
img_name = os.path.splitext(os.path.basename(img_path))[0]
model.test() # test
visuals = model.get_current_visuals(need_HR=need_HR)
img_c = util.tensor2img(visuals['img_c']) # uint8
img_s = util.tensor2img(visuals['img_s']) # uint8
img_p = util.tensor2img(visuals['img_p']) # uint8
# save images
suffix = opt['suffix']
if suffix:
save_c_img_path = os.path.join(dataset_dir, img_name + suffix + '_c.png')
save_s_img_path = os.path.join(dataset_dir, img_name + suffix + '_s.png')
save_p_img_path = os.path.join(dataset_dir, img_name + suffix + '_p.png')
else:
save_c_img_path = os.path.join(dataset_dir, img_name + '_c.png')
save_s_img_path = os.path.join(dataset_dir, img_name + '_s.png')
save_p_img_path = os.path.join(dataset_dir, img_name + '_p.png')
util.save_img(img_c, save_c_img_path)
util.save_img(img_s, save_s_img_path)
util.save_img(img_p, save_p_img_path)
# calculate PSNR and SSIM
if need_HR:
gt_img = util.tensor2img(visuals['HR'])
gt_img = gt_img / 255.
sr_img = img_c / 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)
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,
required=True,
help='Path to options YMAL file.')
opt = option.parse(parser.parse_args().opt, is_train=False)
opt = option.dict_to_nonedict(opt)
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))
util.setup_logger('base',
opt['path']['log'],
'test_' + opt['name'],
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
logger.info(option.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)
logger.info('Number of test images in [{:s}]: {:d}'.format(
dataset_opt['name'], len(test_set)))
test_loaders.append(test_loader)
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 = osp.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_GT = False if test_loader.dataset.opt[
'dataroot_GT'] is None else True
model.feed_data(data, need_GT=need_GT)
img_path = data['GT_path'][0] if need_GT else data['LQ_path'][0]
img_name = osp.splitext(osp.basename(img_path))[0]
model.test()
visuals = model.get_current_visuals(need_GT=need_GT)
sr_img = util.tensor2img(visuals['rlt']) # uint8
# save images
suffix = opt['suffix']
if suffix:
save_img_path = osp.join(dataset_dir,
img_name + suffix + '.png')
else:
save_img_path = osp.join(dataset_dir, img_name + '.png')
util.save_img(sr_img, save_img_path)
# calculate PSNR and SSIM
if need_GT:
gt_img = util.tensor2img(visuals['GT'])
sr_img, gt_img = util.crop_border([sr_img, gt_img],
opt['scale'])
psnr = util.calculate_psnr(sr_img, gt_img)
ssim = util.calculate_ssim(sr_img, gt_img)
test_results['psnr'].append(psnr)
test_results['ssim'].append(ssim)
if gt_img.shape[2] == 3: # RGB image
sr_img_y = bgr2ycbcr(sr_img / 255., only_y=True)
gt_img_y = bgr2ycbcr(gt_img / 255., only_y=True)
psnr_y = util.calculate_psnr(sr_img_y * 255,
gt_img_y * 255)
ssim_y = util.calculate_ssim(sr_img_y * 255,
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)
if need_GT: # 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,
required=True,
help='Path to options JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=False)
util.mkdirs((path for key, path in opt['path'].items()
if not key == 'pretrain_model_G'))
opt = option.dict_to_nonedict(opt)
util.setup_logger(None,
opt['path']['log'],
'test.log',
level=logging.INFO,
screen=True)
logger = logging.getLogger('base')
logger.info(option.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():
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
default='options/train/train_ESRCNN_S2self.json',
help='Path to option JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=True)
opt = option.dict_to_nonedict(opt)
if opt['path']['resume_state']:
resume_state = torch.load(opt['path']['resume_state'])
else:
resume_state = None
util.mkdir_and_rename(opt['path']['experiments_root'])
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))
util.setup_logger(None,
opt['path']['log'],
'train',
level=logging.INFO,
screen=True)
util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO)
logger = logging.getLogger('base')
if resume_state:
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
option.check_resume(opt)
logger.info(option.dict2str(opt))
if opt['use_tb_logger'] and 'debug' not in opt['name']:
from tensorboardX import SummaryWriter
tb_logger = SummaryWriter(log_dir='./tb_logger/' + opt['name'])
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
logger.info('Random seed: {}'.format(seed))
util.set_random_seed(seed)
torch.backends.cudnn.benckmark = True
# Setup TrainDataLoader
trainloader = DataLoader(opt['datasets']['train']['dataroot'],
split='train')
train_size = int(
math.ceil(len(trainloader) / opt['datasets']['train']['batch_size']))
logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
len(trainloader), train_size))
total_iters = int(opt['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
total_epochs, total_iters))
TrainDataLoader = data.DataLoader(
trainloader,
batch_size=opt['datasets']['train']['batch_size'],
num_workers=12,
shuffle=True)
#Setup for validate
valloader = DataLoader(opt['datasets']['train']['dataroot'], split='val')
VALDataLoader = data.DataLoader(
valloader,
batch_size=opt['datasets']['train']['batch_size'] // 5,
num_workers=1,
shuffle=True)
logger.info('Number of val images:{:d}'.format(len(valloader)))
# Setup Model
model = get_model('esrcnn_s2self', opt)
if resume_state:
start_epoch = resume_state['epoch']
current_step = resume_state['iter']
model.resume_training(resume_state)
else:
current_step = 0
start_epoch = 0
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(
start_epoch, current_step))
for epoch in range(start_epoch, total_epochs):
for i, train_data in enumerate(TrainDataLoader):
current_step += 1
if current_step > total_iters:
break
model.update_learning_rate()
model.feed_data(train_data)
model.optimize_parameters(current_step)
if current_step % opt['logger']['print_freq'] == 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[0])
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar(k, v[0], current_step)
logger.info(message)
if current_step % opt['train']['val_freq'] == 0:
avg_psnr = 0.0
idx = 0
for i_val, val_data in enumerate(VALDataLoader):
idx += 1
img_name = val_data[3][0].split('.')[0]
model.feed_data(val_data)
model.val()
visuals = model.get_current_visuals()
pred_img = util.tensor2img(visuals['Pred'])
gt_img = util.tensor2img(visuals['label'])
avg_psnr += util.calculate_psnr(pred_img, gt_img)
avg_psnr = avg_psnr / idx
logger.info('# Validation #PSNR: {:.4e}'.format(avg_psnr))
logger_val = logging.getLogger('val')
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr:{:.4e}'.format(
epoch, current_step, avg_psnr))
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training')
def main():
#### options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
default='options/test/test_KPSAGAN.yml',
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=False)
#### 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:
# 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))
# 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))
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)
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():
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)))
#### create model
model = create_model(opt)
avg_psnr = 0.0
idx = 0
dataset_dir = '/srv/wuyichao/Super-Resolution/KPSAGAN/BasicSR-master/BasicSR-master-c/result_600000/'
util.mkdir(dataset_dir)
for val_data in val_loader:
idx += 1
img_name = os.path.splitext(os.path.basename(
val_data['LQ_path'][0]))[0]
logger.info(img_name)
#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
# save images
suffix = 'cut'
#opt['suffix']
if suffix:
save_img_path = osp.join(dataset_dir, img_name + suffix + '.png')
else:
save_img_path = osp.join(dataset_dir, img_name + '.png')
util.save_img(sr_img, save_img_path)
# calculate PSNR
crop_size = opt['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
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.calculate_psnr(cropped_sr_img * 255,
cropped_gt_img * 255)
avg_psnr = avg_psnr / idx
# log
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
logger_val = logging.getLogger('val') # validation logger
logger_val.info('psnr: {:.4e}'.format(avg_psnr))
def main():
PreUp = False
# 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)
opt = option.dict_to_nonedict(opt) # Convert to NoneDict, which return None for missing key.
ratio = opt["scale"]
if PreUp == True:
ratio=5
# train from scratch OR resume training
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))
# config loggers. Before it, the log will not work
util.setup_logger(None, opt['path']['log'], 'train', level=logging.INFO, screen=True)
util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO)
logger = logging.getLogger('base')
if resume_state:
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
option.check_resume(opt) # check resume options
logger.info(option.dict2str(opt))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
from tensorboardX import SummaryWriter
tb_logger_train = SummaryWriter(log_dir='/mnt/gpid07/users/luis.salgueiro/git/mnt/BasicSR/tb_logger/' + opt['name'] + "/train")
tb_logger_val = SummaryWriter(log_dir='//mnt/gpid07/users/luis.salgueiro/git/mnt/BasicSR/tb_logger/' + opt['name'] + "/val" )
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = 100 #random.randint(1, 10000)
logger.info('Random seed: {}'.format(seed))
util.set_random_seed(seed)
torch.backends.cudnn.benckmark = True
# torch.backends.cudnn.deterministic = True
# print("OLAAAA_-...", os.environ['CUDA_VISIBLE_DEVICES'])
# #########################################
# ######## DATA LOADER ####################
# #########################################
# create train and val dataloader
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
print("Entro DATASET train......")
train_set = create_dataset(dataset_opt)
print("CREO DATASET train_set ", train_set)
train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size']))
logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
total_epochs, total_iters))
train_loader = create_dataloader(train_set, dataset_opt)
print("CREO train loader: ", train_loader)
elif phase == 'val':
print("Entro en phase VAL....")
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
logger.info('Number of val images in [{:s}]: {:d}'.format(dataset_opt['name'],
len(val_set)))
# for _,ii in enumerate(val_loader):
# print("VAL LOADER:........", ii)
# print(val_loader[0])
else:
raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
assert train_loader is not None
assert val_loader is not None
# create model
model = create_model(opt)
#print("PASO..... MODEL ")
# resume training
if resume_state:
print("RESUMING state")
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
print("PASO..... INIT ")
# #########################################
# ######### training ################
# #########################################
# ii=0
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
for epoch in range(start_epoch, total_epochs):
# print("Entro EPOCH...", ii)
for _, train_data in enumerate(train_loader):
# print("Entro TRAIN_LOADER...")
current_step += 1
if current_step > total_iters:
break
# update learning rate
model.update_learning_rate()
# training
#print("....... TRAIN DATA..........", train_data)
model.feed_data(train_data)
model.optimize_parameters(current_step)
# log train
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step, model.get_current_learning_rate())
# print(".............MESSAGE: ", message)
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
#print("MSG: ", message)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
# print("K: ", k)
# print("V: ", v)
if "test" in k:
tb_logger_val.add_scalar(k, v, current_step)
else:
tb_logger_train.add_scalar(k, v, current_step)
logger.info(message)
if current_step % opt['train']['val_freq'] == 0:
avg_psnr_sr = 0.0
avg_psnr_lr = 0.0
avg_psnr_dif = 0.0
avg_ssim_lr, avg_ssim_sr, avg_ssim_dif = 0.0, 0.0, 0.0
avg_ergas_lr, avg_ergas_sr, avg_ergas_dif = 0.0, 0.0, 0.0
idx = 0
# for val_data in val_loader:
for _, val_data in enumerate(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)
# print("Img nameVaL: ", img_name)
model.feed_data(val_data)
model.test()
visuals = model.get_current_visuals()
sr_img = util.tensor2imgNorm(visuals['SR'],out_type=np.uint8, min_max=(0, 1), MinVal=val_data["LR_min"], MaxVal=val_data["LR_max"]) # uint16
gt_img = util.tensor2imgNorm(visuals['HR'],out_type=np.uint8, min_max=(0, 1), MinVal=val_data["HR_min"], MaxVal=val_data["HR_max"]) # uint16
lr_img = util.tensor2imgNorm(visuals['LR'], out_type=np.uint8, min_max=(0, 1),
MinVal=val_data["LR_min"], MaxVal=val_data["LR_max"]) # uint16
# Save SR images for reference
if idx < 10:
# print(idx)
util.mkdir(img_dir)
save_img_path = os.path.join(img_dir, '{:s}_{:d}'.format(img_name, current_step))
util.save_imgSR(sr_img, save_img_path)
util.save_imgHR(gt_img, save_img_path)
util.save_imgLR(lr_img, save_img_path)
print("SAVING CROPS")
util.save_imgCROP(lr_img,gt_img,sr_img , save_img_path, ratio, PreUp=PreUp)
if PreUp==False:
dim2 = (gt_img.shape[1], gt_img.shape[1])
print("DIM:", dim2)
print("LR image shape ", lr_img.shape)
print("HR image shape ", gt_img.shape)
lr_img = cv2.resize(np.transpose(lr_img,(1,2,0)), dim2, interpolation=cv2.INTER_NEAREST)
lr_img = np.transpose(lr_img,(2,0,1))
print("LR image 2 shape ", lr_img.shape)
print("LR image 2 shape ", lr_img.shape)
avg_psnr_sr += util.calculate_psnr2(sr_img, gt_img)
avg_psnr_lr += util.calculate_psnr2(lr_img, gt_img)
avg_ssim_lr += util.calculate_ssim2(lr_img, gt_img)
avg_ssim_sr += util.calculate_ssim2(sr_img, gt_img)
avg_ergas_lr += util.calculate_ergas(lr_img, gt_img, pixratio=ratio)
avg_ergas_sr += util.calculate_ergas(sr_img, gt_img, pixratio=ratio)
#avg_psnr += util.calculate_psnr2(cropped_sr_img, cropped_gt_img)
avg_psnr_sr = avg_psnr_sr / idx
avg_psnr_lr = avg_psnr_lr / idx
avg_psnr_dif = avg_psnr_lr - avg_psnr_sr
avg_ssim_lr = avg_ssim_lr / idx
avg_ssim_sr = avg_ssim_sr / idx
avg_ssim_dif = avg_ssim_lr - avg_ssim_sr
avg_ergas_lr = avg_ergas_lr / idx
avg_ergas_sr = avg_ergas_sr / idx
avg_ergas_dif = avg_ergas_lr - avg_ergas_sr
# print("IDX: ", idx)
# log VALIDATION
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr_sr))
logger.info('# Validation # SSIM: {:.4e}'.format(avg_ssim_sr))
logger.info('# Validation # ERGAS: {:.4e}'.format(avg_ergas_sr))
logger_val = logging.getLogger('val') # validation logger
logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr_SR: {:.4e}'.format(
epoch, current_step, avg_psnr_sr))
logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr_LR: {:.4e}'.format(
epoch, current_step, avg_psnr_lr))
logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr_DIF: {:.4e}'.format(
epoch, current_step, avg_psnr_dif))
logger_val.info('<epoch:{:3d}, iter:{:8,d}> ssim_LR: {:.4e}'.format(
epoch, current_step, avg_ssim_lr))
logger_val.info('<epoch:{:3d}, iter:{:8,d}> ssim_SR: {:.4e}'.format(
epoch, current_step, avg_ssim_sr))
logger_val.info('<epoch:{:3d}, iter:{:8,d}> ssim_DIF: {:.4e}'.format(
epoch, current_step, avg_ssim_dif))
logger_val.info('<epoch:{:3d}, iter:{:8,d}> ergas_LR: {:.4e}'.format(
epoch, current_step, avg_ergas_lr))
logger_val.info('<epoch:{:3d}, iter:{:8,d}> ergas_SR: {:.4e}'.format(
epoch, current_step, avg_ergas_sr))
logger_val.info('<epoch:{:3d}, iter:{:8,d}> ergas_DIF: {:.4e}'.format(
epoch, current_step, avg_ergas_dif))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger_val.add_scalar('dif_PSNR', avg_psnr_dif, current_step)
# tb_logger.add_scalar('psnr', avg_psnr, current_step)
tb_logger_val.add_scalar('dif_SSIM', avg_ssim_dif, current_step)
tb_logger_val.add_scalar('dif_ERGAS', avg_ergas_dif, current_step)
tb_logger_val.add_scalar('psnr_LR', avg_psnr_lr, current_step)
# tb_logger.add_scalar('psnr', avg_psnr, current_step)
tb_logger_val.add_scalar('ssim_LR', avg_ssim_lr, current_step)
tb_logger_val.add_scalar('ERGAS_LR', avg_ergas_lr, current_step)
tb_logger_val.add_scalar('psnr_SR', avg_psnr_sr, current_step)
# tb_logger.add_scalar('psnr', avg_psnr, current_step)
tb_logger_val.add_scalar('ssim_SR', avg_ssim_sr, current_step)
tb_logger_val.add_scalar('ERGAS_SR', avg_ergas_sr, current_step)
print("****** SR_IMG: ", sr_img.shape)
print("****** LR_IMG: ", lr_img.shape)
print("****** GT_IMG: ", gt_img.shape)
fig1,ax1 = ep.plot_rgb(sr_img, rgb=[2, 1, 0], stretch=True)
tb_logger_val.add_figure("SR_plt", fig1, current_step,close=True)
fig2, ax2 = ep.plot_rgb(gt_img, rgb=[2, 1, 0], stretch=True)
tb_logger_val.add_figure("GT_plt", fig2, current_step, close=True)
fig3, ax3 = ep.plot_rgb(lr_img, rgb=[2, 1, 0], stretch=True)
tb_logger_val.add_figure("LR_plt", fig3, current_step, close=True)
# print("TERMINO GUARDAR IMG TB")
# save models and training states
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
# ii=ii+1
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training.')
def main():
############################################
#
# set options
#
############################################
parser = argparse.ArgumentParser()
parser.add_argument('--opt', type=str, help='Path to option YAML 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()
print("Rank:", rank)
print("------------------DIST-------------------------")
############################################
#
# 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:
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))
# 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('base_val',
opt['path']['log'],
'val_' + opt['name'],
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
logger_val = logging.getLogger('base_val')
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
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
else:
# config loggers. Before it, the log will not work
util.setup_logger('base',
opt['path']['log'],
'train_',
level=logging.INFO,
screen=True)
print("set train log")
util.setup_logger('base_val',
opt['path']['log'],
'val_',
level=logging.INFO,
screen=True)
print("set val log")
logger = logging.getLogger('base')
logger_val = logging.getLogger('base_val')
# convert to NoneDict, which returns None for missing keys
opt = option.dict_to_nonedict(opt)
#### 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.benchmark = True
# torch.backends.cudnn.deterministic = True
############################################
#
# create train and val dataloader
#
############################################
####
# dataset_ratio = 200 # enlarge the size of each epoch, todo: what it is
dataset_ratio = 1 # enlarge the size of each epoch, todo: what it is
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))
total_iters = train_size
total_epochs = int(opt['train']['epoch'])
if opt['dist']:
train_sampler = DistIterSampler(train_set, world_size, rank,
dataset_ratio)
# total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio)))
total_epochs = int(opt['train']['epoch'])
if opt['train']['enable'] == False:
total_epochs = 1
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 = create_model(opt)
print("Model Created! ")
#### 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
print("Not Resume Training")
############################################
#
# training
#
############################################
####
####
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(
start_epoch, current_step))
Avg_train_loss = AverageMeter() # total
if (opt['train']['pixel_criterion'] == 'cb+ssim'):
Avg_train_loss_pix = AverageMeter()
Avg_train_loss_ssim = AverageMeter()
elif (opt['train']['pixel_criterion'] == 'cb+ssim+vmaf'):
Avg_train_loss_pix = AverageMeter()
Avg_train_loss_ssim = AverageMeter()
Avg_train_loss_vmaf = AverageMeter()
elif (opt['train']['pixel_criterion'] == 'ssim'):
Avg_train_loss_ssim = AverageMeter()
elif (opt['train']['pixel_criterion'] == 'msssim'):
Avg_train_loss_msssim = AverageMeter()
elif (opt['train']['pixel_criterion'] == 'cb+msssim'):
Avg_train_loss_pix = AverageMeter()
Avg_train_loss_msssim = AverageMeter()
saved_total_loss = 10e10
saved_total_PSNR = -1
for epoch in range(start_epoch, total_epochs):
############################################
#
# Start a new epoch
#
############################################
# Turn into training mode
#model = model.train()
# reset total loss
Avg_train_loss.reset()
current_step = 0
if (opt['train']['pixel_criterion'] == 'cb+ssim'):
Avg_train_loss_pix.reset()
Avg_train_loss_ssim.reset()
elif (opt['train']['pixel_criterion'] == 'cb+ssim+vmaf'):
Avg_train_loss_pix.reset()
Avg_train_loss_ssim.reset()
Avg_train_loss_vmaf.reset()
elif (opt['train']['pixel_criterion'] == 'ssim'):
Avg_train_loss_ssim = AverageMeter()
elif (opt['train']['pixel_criterion'] == 'msssim'):
Avg_train_loss_msssim = AverageMeter()
elif (opt['train']['pixel_criterion'] == 'cb+msssim'):
Avg_train_loss_pix = AverageMeter()
Avg_train_loss_msssim = AverageMeter()
if opt['dist']:
train_sampler.set_epoch(epoch)
for train_idx, train_data in enumerate(train_loader):
if 'debug' in opt['name']:
img_dir = os.path.join(opt['path']['train_images'])
util.mkdir(img_dir)
LQ = train_data['LQs']
GT = train_data['GT']
GT_img = util.tensor2img(GT) # uint8
save_img_path = os.path.join(
img_dir, '{:4d}_{:s}.png'.format(train_idx, 'debug_GT'))
util.save_img(GT_img, save_img_path)
for i in range(5):
LQ_img = util.tensor2img(LQ[0, i, ...]) # uint8
save_img_path = os.path.join(
img_dir,
'{:4d}_{:s}_{:1d}.png'.format(train_idx, 'debug_LQ',
i))
util.save_img(LQ_img, save_img_path)
if (train_idx >= 3):
break
if opt['train']['enable'] == False:
message_train_loss = 'None'
break
current_step += 1
if current_step > total_iters:
print("Total Iteration Reached !")
break
#### update learning rate
if opt['train']['lr_scheme'] == 'ReduceLROnPlateau':
pass
else:
model.update_learning_rate(
current_step, warmup_iter=opt['train']['warmup_iter'])
#### training
model.feed_data(train_data)
# if opt['train']['lr_scheme'] == 'ReduceLROnPlateau':
# model.optimize_parameters_without_schudlue(current_step)
# else:
model.optimize_parameters(current_step)
if (opt['train']['pixel_criterion'] == 'cb+ssim'):
Avg_train_loss.update(model.log_dict['total_loss'], 1)
Avg_train_loss_pix.update(model.log_dict['l_pix'], 1)
Avg_train_loss_ssim.update(model.log_dict['ssim_loss'], 1)
elif (opt['train']['pixel_criterion'] == 'cb+ssim+vmaf'):
Avg_train_loss.update(model.log_dict['total_loss'], 1)
Avg_train_loss_pix.update(model.log_dict['l_pix'], 1)
Avg_train_loss_ssim.update(model.log_dict['ssim_loss'], 1)
Avg_train_loss_vmaf.update(model.log_dict['vmaf_loss'], 1)
elif (opt['train']['pixel_criterion'] == 'ssim'):
Avg_train_loss.update(model.log_dict['total_loss'], 1)
Avg_train_loss_ssim.update(model.log_dict['ssim_loss'], 1)
elif (opt['train']['pixel_criterion'] == 'msssim'):
Avg_train_loss.update(model.log_dict['total_loss'], 1)
Avg_train_loss_msssim.update(model.log_dict['msssim_loss'], 1)
elif (opt['train']['pixel_criterion'] == 'cb+msssim'):
Avg_train_loss.update(model.log_dict['total_loss'], 1)
Avg_train_loss_pix.update(model.log_dict['l_pix'], 1)
Avg_train_loss_msssim.update(model.log_dict['msssim_loss'], 1)
else:
Avg_train_loss.update(model.log_dict['l_pix'], 1)
# add total train loss
if (opt['train']['pixel_criterion'] == 'cb+ssim'):
message_train_loss = ' pix_avg_loss: {:.4e}'.format(
Avg_train_loss_pix.avg)
message_train_loss += ' ssim_avg_loss: {:.4e}'.format(
Avg_train_loss_ssim.avg)
message_train_loss += ' total_avg_loss: {:.4e}'.format(
Avg_train_loss.avg)
elif (opt['train']['pixel_criterion'] == 'cb+ssim+vmaf'):
message_train_loss = ' pix_avg_loss: {:.4e}'.format(
Avg_train_loss_pix.avg)
message_train_loss += ' ssim_avg_loss: {:.4e}'.format(
Avg_train_loss_ssim.avg)
message_train_loss += ' vmaf_avg_loss: {:.4e}'.format(
Avg_train_loss_vmaf.avg)
message_train_loss += ' total_avg_loss: {:.4e}'.format(
Avg_train_loss.avg)
elif (opt['train']['pixel_criterion'] == 'ssim'):
message_train_loss = ' ssim_avg_loss: {:.4e}'.format(
Avg_train_loss_ssim.avg)
message_train_loss += ' total_avg_loss: {:.4e}'.format(
Avg_train_loss.avg)
elif (opt['train']['pixel_criterion'] == 'msssim'):
message_train_loss = ' msssim_avg_loss: {:.4e}'.format(
Avg_train_loss_msssim.avg)
message_train_loss += ' total_avg_loss: {:.4e}'.format(
Avg_train_loss.avg)
elif (opt['train']['pixel_criterion'] == 'cb+msssim'):
message_train_loss = ' pix_avg_loss: {:.4e}'.format(
Avg_train_loss_pix.avg)
message_train_loss += ' msssim_avg_loss: {:.4e}'.format(
Avg_train_loss_msssim.avg)
message_train_loss += ' total_avg_loss: {:.4e}'.format(
Avg_train_loss.avg)
else:
message_train_loss = ' train_avg_loss: {:.4e}'.format(
Avg_train_loss.avg)
#### log
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
message = '[epoch:{:3d}, iter:{:8,d}, lr:('.format(
epoch, current_step)
for v in model.get_current_learning_rate():
message += '{:.3e},'.format(v)
message += ')] '
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
message += message_train_loss
if rank <= 0:
logger.info(message)
############################################
#
# end of one epoch, save epoch model
#
############################################
#### save models and training states
# if current_step % opt['logger']['save_checkpoint_freq'] == 0:
# if rank <= 0:
# logger.info('Saving models and training states.')
# model.save(current_step)
# model.save('latest')
# # model.save_training_state(epoch, current_step)
# # todo delete previous weights
# previous_step = current_step - opt['logger']['save_checkpoint_freq']
# save_filename = '{}_{}.pth'.format(previous_step, 'G')
# save_path = os.path.join(opt['path']['models'], save_filename)
# if os.path.exists(save_path):
# os.remove(save_path)
if epoch == 1:
save_filename = '{:04d}_{}.pth'.format(0, 'G')
save_path = os.path.join(opt['path']['models'], save_filename)
if os.path.exists(save_path):
os.remove(save_path)
save_filename = '{:04d}_{}.pth'.format(epoch - 1, 'G')
save_path = os.path.join(opt['path']['models'], save_filename)
if os.path.exists(save_path):
os.remove(save_path)
if rank <= 0:
logger.info('Saving models and training states.')
save_filename = '{:04d}'.format(epoch)
model.save(save_filename)
# model.save('latest')
# model.save_training_state(epoch, current_step)
############################################
#
# end of one epoch, do validation
#
############################################
#### validation
#if opt['datasets'].get('val', None) and current_step % opt['train']['val_freq'] == 0:
if opt['datasets'].get('val', None):
if opt['model'] in [
'sr', 'srgan'
] and rank <= 0: # image restoration validation
# does not support multi-GPU validation
pbar = util.ProgressBar(len(val_loader))
avg_psnr = 0.
idx = 0
for val_data in val_loader:
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['rlt']) # uint8
gt_img = util.tensor2img(visuals['GT']) # 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
sr_img, gt_img = util.crop_border([sr_img, gt_img],
opt['scale'])
avg_psnr += util.calculate_psnr(sr_img, gt_img)
pbar.update('Test {}'.format(img_name))
avg_psnr = avg_psnr / idx
# log
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
else: # video restoration validation
if opt['dist']:
# todo : multi-GPU testing
psnr_rlt = {} # with border and center frames
psnr_rlt_avg = {}
psnr_total_avg = 0.
ssim_rlt = {} # with border and center frames
ssim_rlt_avg = {}
ssim_total_avg = 0.
val_loss_rlt = {}
val_loss_rlt_avg = {}
val_loss_total_avg = 0.
if rank == 0:
pbar = util.ProgressBar(len(val_set))
for idx in range(rank, len(val_set), world_size):
print('idx', idx)
if 'debug' in opt['name']:
if (idx >= 3):
break
val_data = val_set[idx]
val_data['LQs'].unsqueeze_(0)
val_data['GT'].unsqueeze_(0)
folder = val_data['folder']
idx_d, max_idx = val_data['idx'].split('/')
idx_d, max_idx = int(idx_d), int(max_idx)
if psnr_rlt.get(folder, None) is None:
psnr_rlt[folder] = torch.zeros(max_idx,
dtype=torch.float32,
device='cuda')
if ssim_rlt.get(folder, None) is None:
ssim_rlt[folder] = torch.zeros(max_idx,
dtype=torch.float32,
device='cuda')
if val_loss_rlt.get(folder, None) is None:
val_loss_rlt[folder] = torch.zeros(
max_idx, dtype=torch.float32, device='cuda')
# tmp = torch.zeros(max_idx, dtype=torch.float32, device='cuda')
model.feed_data(val_data)
# model.test()
# model.test_stitch()
if opt['stitch'] == True:
model.test_stitch()
else:
model.test() # large GPU memory
# visuals = model.get_current_visuals()
visuals = model.get_current_visuals(
save=True,
name='{}_{}'.format(folder, idx),
save_path=opt['path']['val_images'])
rlt_img = util.tensor2img(visuals['rlt']) # uint8
gt_img = util.tensor2img(visuals['GT']) # uint8
# calculate PSNR
psnr = util.calculate_psnr(rlt_img, gt_img)
psnr_rlt[folder][idx_d] = psnr
# calculate SSIM
ssim = util.calculate_ssim(rlt_img, gt_img)
ssim_rlt[folder][idx_d] = ssim
# calculate Val loss
val_loss = model.get_loss()
val_loss_rlt[folder][idx_d] = val_loss
logger.info(
'{}_{:02d} PSNR: {:.4f}, SSIM: {:.4f}'.format(
folder, idx, psnr, ssim))
if rank == 0:
for _ in range(world_size):
pbar.update('Test {} - {}/{}'.format(
folder, idx_d, max_idx))
# # collect data
for _, v in psnr_rlt.items():
dist.reduce(v, 0)
for _, v in ssim_rlt.items():
dist.reduce(v, 0)
for _, v in val_loss_rlt.items():
dist.reduce(v, 0)
dist.barrier()
if rank == 0:
psnr_rlt_avg = {}
psnr_total_avg = 0.
for k, v in psnr_rlt.items():
psnr_rlt_avg[k] = torch.mean(v).cpu().item()
psnr_total_avg += psnr_rlt_avg[k]
psnr_total_avg /= len(psnr_rlt)
log_s = '# Validation # PSNR: {:.4e}:'.format(
psnr_total_avg)
for k, v in psnr_rlt_avg.items():
log_s += ' {}: {:.4e}'.format(k, v)
logger.info(log_s)
# ssim
ssim_rlt_avg = {}
ssim_total_avg = 0.
for k, v in ssim_rlt.items():
ssim_rlt_avg[k] = torch.mean(v).cpu().item()
ssim_total_avg += ssim_rlt_avg[k]
ssim_total_avg /= len(ssim_rlt)
log_s = '# Validation # PSNR: {:.4e}:'.format(
ssim_total_avg)
for k, v in ssim_rlt_avg.items():
log_s += ' {}: {:.4e}'.format(k, v)
logger.info(log_s)
# added
val_loss_rlt_avg = {}
val_loss_total_avg = 0.
for k, v in val_loss_rlt.items():
val_loss_rlt_avg[k] = torch.mean(v).cpu().item()
val_loss_total_avg += val_loss_rlt_avg[k]
val_loss_total_avg /= len(val_loss_rlt)
log_l = '# Validation # Loss: {:.4e}:'.format(
val_loss_total_avg)
for k, v in val_loss_rlt_avg.items():
log_l += ' {}: {:.4e}'.format(k, v)
logger.info(log_l)
message = ''
for v in model.get_current_learning_rate():
message += '{:.5e}'.format(v)
logger_val.info(
'Epoch {:02d}, LR {:s}, PSNR {:.4f}, SSIM {:.4f} Train {:s}, Val Total Loss {:.4e}'
.format(epoch, message, psnr_total_avg,
ssim_total_avg, message_train_loss,
val_loss_total_avg))
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr_avg', psnr_total_avg,
current_step)
for k, v in psnr_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
# add val loss
tb_logger.add_scalar('val_loss_avg',
val_loss_total_avg,
current_step)
for k, v in val_loss_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
else: # Todo: our function One GPU
pbar = util.ProgressBar(len(val_loader))
psnr_rlt = {} # with border and center frames
psnr_rlt_avg = {}
psnr_total_avg = 0.
ssim_rlt = {} # with border and center frames
ssim_rlt_avg = {}
ssim_total_avg = 0.
val_loss_rlt = {}
val_loss_rlt_avg = {}
val_loss_total_avg = 0.
for val_inx, val_data in enumerate(val_loader):
if 'debug' in opt['name']:
if (val_inx >= 5):
break
folder = val_data['folder'][0]
# idx_d = val_data['idx'].item()
idx_d = val_data['idx']
# border = val_data['border'].item()
if psnr_rlt.get(folder, None) is None:
psnr_rlt[folder] = []
if ssim_rlt.get(folder, None) is None:
ssim_rlt[folder] = []
if val_loss_rlt.get(folder, None) is None:
val_loss_rlt[folder] = []
# process the black blank [B N C H W]
print(val_data['LQs'].size())
H_S = val_data['LQs'].size(3) # 540
W_S = val_data['LQs'].size(4) # 960
print(H_S)
print(W_S)
blank_1_S = 0
blank_2_S = 0
print(val_data['LQs'][0, 2, 0, :, :].size())
for i in range(H_S):
if not sum(val_data['LQs'][0, 2, 0, i, :]) == 0:
blank_1_S = i - 1
# assert not sum(data_S[:, :, 0][i+1]) == 0
break
for i in range(H_S):
if not sum(val_data['LQs'][0, 2, 0, :,
H_S - i - 1]) == 0:
blank_2_S = (H_S - 1) - i - 1
# assert not sum(data_S[:, :, 0][blank_2_S-1]) == 0
break
print('LQ :', blank_1_S, blank_2_S)
if blank_1_S == -1:
print('LQ has no blank')
blank_1_S = 0
blank_2_S = H_S
# val_data['LQs'] = val_data['LQs'][:,:,:,blank_1_S:blank_2_S,:]
print("LQ", val_data['LQs'].size())
# end of process the black blank
model.feed_data(val_data)
if opt['stitch'] == True:
model.test_stitch()
else:
model.test() # large GPU memory
# process blank
blank_1_L = blank_1_S << 2
blank_2_L = blank_2_S << 2
print(blank_1_L, blank_2_L)
print(model.fake_H.size())
if not blank_1_S == 0:
# model.fake_H = model.fake_H[:,:,blank_1_L:blank_2_L,:]
model.fake_H[:, :, 0:blank_1_L, :] = 0
model.fake_H[:, :, blank_2_L:H_S, :] = 0
# end of # process blank
visuals = model.get_current_visuals(
save=True,
name='{}_{:02d}'.format(folder, val_inx),
save_path=opt['path']['val_images'])
rlt_img = util.tensor2img(visuals['rlt']) # uint8
gt_img = util.tensor2img(visuals['GT']) # uint8
# calculate PSNR
psnr = util.calculate_psnr(rlt_img, gt_img)
psnr_rlt[folder].append(psnr)
# calculate SSIM
ssim = util.calculate_ssim(rlt_img, gt_img)
ssim_rlt[folder].append(ssim)
# val loss
val_loss = model.get_loss()
val_loss_rlt[folder].append(val_loss.item())
logger.info(
'{}_{:02d} PSNR: {:.4f}, SSIM: {:.4f}'.format(
folder, val_inx, psnr, ssim))
pbar.update('Test {} - {}'.format(folder, idx_d))
# average PSNR
for k, v in psnr_rlt.items():
psnr_rlt_avg[k] = sum(v) / len(v)
psnr_total_avg += psnr_rlt_avg[k]
psnr_total_avg /= len(psnr_rlt)
log_s = '# Validation # PSNR: {:.4e}:'.format(
psnr_total_avg)
for k, v in psnr_rlt_avg.items():
log_s += ' {}: {:.4e}'.format(k, v)
logger.info(log_s)
# average SSIM
for k, v in ssim_rlt.items():
ssim_rlt_avg[k] = sum(v) / len(v)
ssim_total_avg += ssim_rlt_avg[k]
ssim_total_avg /= len(ssim_rlt)
log_s = '# Validation # SSIM: {:.4e}:'.format(
ssim_total_avg)
for k, v in ssim_rlt_avg.items():
log_s += ' {}: {:.4e}'.format(k, v)
logger.info(log_s)
# average VMAF
# average Val LOSS
for k, v in val_loss_rlt.items():
val_loss_rlt_avg[k] = sum(v) / len(v)
val_loss_total_avg += val_loss_rlt_avg[k]
val_loss_total_avg /= len(val_loss_rlt)
log_l = '# Validation # Loss: {:.4e}:'.format(
val_loss_total_avg)
for k, v in val_loss_rlt_avg.items():
log_l += ' {}: {:.4e}'.format(k, v)
logger.info(log_l)
# toal validation log
message = ''
for v in model.get_current_learning_rate():
message += '{:.5e}'.format(v)
logger_val.info(
'Epoch {:02d}, LR {:s}, PSNR {:.4f}, SSIM {:.4f} Train {:s}, Val Total Loss {:.4e}'
.format(epoch, message, psnr_total_avg, ssim_total_avg,
message_train_loss, val_loss_total_avg))
# end add
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr_avg', psnr_total_avg,
current_step)
for k, v in psnr_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
# tb_logger.add_scalar('ssim_avg', ssim_total_avg, current_step)
# for k, v in ssim_rlt_avg.items():
# tb_logger.add_scalar(k, v, current_step)
# add val loss
tb_logger.add_scalar('val_loss_avg',
val_loss_total_avg, current_step)
for k, v in val_loss_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
############################################
#
# end of validation, save model
#
############################################
#
logger.info("Finished an epoch, Check and Save the model weights")
# we check the validation loss instead of training loss. OK~
if saved_total_loss >= val_loss_total_avg:
saved_total_loss = val_loss_total_avg
#torch.save(model.state_dict(), args.save_path + "/best" + ".pth")
model.save('best')
logger.info(
"Best Weights updated for decreased validation loss")
else:
logger.info(
"Weights Not updated for undecreased validation loss")
if saved_total_PSNR <= psnr_total_avg:
saved_total_PSNR = psnr_total_avg
model.save('bestPSNR')
logger.info(
"Best Weights updated for increased validation PSNR")
else:
logger.info(
"Weights Not updated for unincreased validation PSNR")
############################################
#
# end of one epoch, schedule LR
#
############################################
# add scheduler todo
if opt['train']['lr_scheme'] == 'ReduceLROnPlateau':
for scheduler in model.schedulers:
# scheduler.step(val_loss_total_avg)
scheduler.step(val_loss_total_avg)
if rank <= 0:
logger.info('Saving the final model.')
model.save('last')
logger.info('End of training.')
tb_logger.close()
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
opt['dist'] = False
rank = -1
print('Disabled distributed training.')
#### loading resume state if exists
if opt['path'].get('resume_state', None):
resume_state_path, _ = get_resume_paths(opt)
print('\n\t************************')
print('resume_state_path: ', resume_state_path)
print('\n\t************************')
# distributed resuming: all load into default GPU
if resume_state_path is None:
resume_state = None
else:
device_id = torch.cuda.current_device()
resume_state = torch.load(
resume_state_path,
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:
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))
# 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.get('use_tb_logger', False) 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
conf_name = basename(args.opt).replace(".yml", "")
exp_dir = opt['path']['experiments_root']
log_dir_train = os.path.join(exp_dir, 'tb', conf_name, 'train')
log_dir_valid = os.path.join(exp_dir, 'tb', conf_name, 'valid')
tb_logger_train = SummaryWriter(log_dir=log_dir_train)
tb_logger_valid = SummaryWriter(log_dir=log_dir_valid)
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)
#### 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.benchmark = 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)
print('Dataset created')
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))
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
current_step = 0 if resume_state is None else resume_state['iter']
model = create_model(opt, current_step)
#### 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
timer = Timer()
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(
start_epoch, current_step))
timerData = TickTock()
for epoch in range(start_epoch, total_epochs + 1):
if opt['dist']:
train_sampler.set_epoch(epoch)
timerData.tick()
for _, train_data in enumerate(train_loader):
timerData.tock()
current_step += 1
if current_step > total_iters:
break
#### training
model.feed_data(train_data)
#### update learning rate
model.update_learning_rate(current_step,
warmup_iter=opt['train']['warmup_iter'])
nll = model.optimize_parameters(current_step)
"""
try:
nll = model.optimize_parameters(current_step)
except RuntimeError as e:
nll = None
print("Skipping ERROR caught in nll = model.optimize_parameters(current_step): ")
print(e)
"""
if nll is None:
nll = 0
#### log
def eta(t_iter):
return (t_iter * (opt['train']['niter'] - current_step)) / 3600
if current_step % opt['logger']['print_freq'] == 0 \
or current_step - (resume_state['iter'] if resume_state else 0) < 25:
avg_time = timer.get_average_and_reset()
avg_data_time = timerData.get_average_and_reset()
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}, t:{:.2e}, td:{:.2e}, eta:{:.2e}, nll:{:.3e}> '.format(
epoch, current_step, model.get_current_learning_rate(),
avg_time, avg_data_time, eta(avg_time), nll)
print(message)
timer.tick()
# Reduce number of logs
if current_step % 5 == 0:
tb_logger_train.add_scalar('loss/nll', nll, current_step)
tb_logger_train.add_scalar('lr/base',
model.get_current_learning_rate(),
current_step)
tb_logger_train.add_scalar('time/iteration',
timer.get_last_iteration(),
current_step)
tb_logger_train.add_scalar('time/data',
timerData.get_last_iteration(),
current_step)
tb_logger_train.add_scalar('time/eta',
eta(timer.get_last_iteration()),
current_step)
for k, v in model.get_current_log().items():
tb_logger_train.add_scalar(k, v, current_step)
# validation
if current_step % opt['train']['val_freq'] == 0 and rank <= 0:
avg_psnr = 0.0
idx = 0
nlls = []
for val_data in val_loader:
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)
nll = model.test()
if nll is None:
nll = 0
nlls.append(nll)
visuals = model.get_current_visuals()
sr_img = None
# Save SR images for reference
if hasattr(model, 'heats'):
for heat in model.heats:
for i in range(model.n_sample):
sr_img = util.tensor2img(
visuals['SR', heat, i]) # uint8
save_img_path = os.path.join(
img_dir,
'{:s}_{:09d}_h{:03d}_s{:d}.png'.format(
img_name, current_step,
int(heat * 100), i))
util.save_img(sr_img, save_img_path)
else:
sr_img = util.tensor2img(visuals['SR']) # uint8
save_img_path = os.path.join(
img_dir,
'{:s}_{:d}.png'.format(img_name, current_step))
util.save_img(sr_img, save_img_path)
assert sr_img is not None
# Save LQ images for reference
save_img_path_lq = os.path.join(
img_dir, '{:s}_LQ.png'.format(img_name))
if not os.path.isfile(save_img_path_lq):
lq_img = util.tensor2img(visuals['LQ']) # uint8
util.save_img(
cv2.resize(lq_img,
dsize=None,
fx=opt['scale'],
fy=opt['scale'],
interpolation=cv2.INTER_NEAREST),
save_img_path_lq)
# Save GT images for reference
gt_img = util.tensor2img(visuals['GT']) # uint8
save_img_path_gt = os.path.join(
img_dir, '{:s}_GT.png'.format(img_name))
if not os.path.isfile(save_img_path_gt):
util.save_img(gt_img, save_img_path_gt)
# calculate PSNR
crop_size = opt['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
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.calculate_psnr(cropped_sr_img * 255,
cropped_gt_img * 255)
avg_psnr = avg_psnr / idx
avg_nll = sum(nlls) / len(nlls)
# log
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
logger_val = logging.getLogger('val') # validation logger
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}'.format(
epoch, current_step, avg_psnr))
# tensorboard logger
tb_logger_valid.add_scalar('loss/psnr', avg_psnr, current_step)
tb_logger_valid.add_scalar('loss/nll', avg_nll, current_step)
tb_logger_train.flush()
tb_logger_valid.flush()
#### save models and training states
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
if rank <= 0:
logger.info('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
timerData.tick()
with open(os.path.join(opt['path']['root'], "TRAIN_DONE"), 'w') as f:
f.write("TRAIN_DONE")
if rank <= 0:
logger.info('Saving the final model.')
model.save('latest')
logger.info('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(), is_train=True)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# train from scratch OR resume training
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))
# config loggers. Before it, the log will not work
util.setup_logger(None,
opt['path']['log'],
'train',
level=logging.INFO,
screen=True)
util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO)
logger = logging.getLogger('base')
if resume_state:
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
option.check_resume(opt) # check resume options
logger.info(option.dict2str(opt))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
from tensorboardX import SummaryWriter
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
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
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']))
logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
len(train_set), train_size))
total_iters = int(opt['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
total_epochs, 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)
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 = create_model(opt)
# resume training
if resume_state:
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):
for _, train_data in enumerate(train_loader):
current_step += 1
if current_step > total_iters:
break
# validation
if current_step % opt['train']['val_freq'] == 0:
visuals = model.get_current_visuals(isTrain=True)
os.makedirs('valid/' + opt['name'], exist_ok=True)
util.save_img(
util.tensor2img(visuals['LR']), 'valid/' + opt['name'] +
'/' + str(current_step) + '_LR.png')
util.save_img(
util.tensor2img(visuals['HR']), 'valid/' + opt['name'] +
'/' + str(current_step) + '_HR.png')
util.save_img(
util.tensor2img(visuals['SR']), 'valid/' + opt['name'] +
'/' + str(current_step) + '_SR.png')
util.save_img(
util.tensor2img(visuals['SRgray']), 'valid/' +
opt['name'] + '/' + str(current_step) + '_SRgray.png')
util.save_img(
util.tensor2img(visuals['HRgray']), 'valid/' +
opt['name'] + '/' + str(current_step) + '_HRgray.png')
# update learning rate
model.update_learning_rate()
# training
model.feed_data(train_data)
model.optimize_parameters(current_step)
# log
if current_step % opt['logger']['print_freq'] == 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)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar(k, v, current_step)
logger.info(message)
# save models and training states
if current_step % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(current_step)
model.save_training_state(epoch, current_step)
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training.')
def main():
#### options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, help='Path to option YAML 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:
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 'wandb_load_run_path' not in key))
# 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)
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
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
if opt['use_wandb_logger'] and 'debug' not in opt['name']:
json_path = os.path.join(os.path.expanduser('~'),
'.wandb_api_keys.json')
if os.path.exists(json_path):
with open(json_path, 'r') as j:
json_file = json.loads(j.read())
os.environ['WANDB_API_KEY'] = json_file['ryul99']
wandb.init(project="mmsr", config=opt, sync_tensorboard=True)
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)
#### 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))
if opt['use_wandb_logger'] and 'debug' not in opt['name']:
wandb.config.update({'random_seed': seed})
util.set_random_seed(seed)
torch.backends.cudnn.benchmark = 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 = create_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)
for _, train_data in enumerate(train_loader):
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,
noise_mode=opt['datasets']['train']['noise_mode'],
noise_rate=opt['datasets']['train']['noise_rate'])
model.optimize_parameters(current_step)
#### log
if current_step % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
message = '[epoch:{:3d}, iter:{:8,d}, lr:('.format(
epoch, current_step)
for v in model.get_current_learning_rate():
message += '{:.3e},'.format(v)
message += ')] '
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if opt['use_wandb_logger'] and 'debug' not in opt['name']:
if rank <= 0:
wandb.log({k: v}, step=current_step)
if rank <= 0:
logger.info(message)
#### validation
if opt['datasets'].get(
'val',
None) and current_step % opt['train']['val_freq'] == 0:
if opt['model'] in [
'sr', 'srgan'
] and rank <= 0: # image restoration validation
# does not support multi-GPU validation
pbar = util.ProgressBar(len(val_loader))
avg_psnr = 0.
idx = 0
for val_data in val_loader:
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,
noise_mode=opt['datasets']['val']['noise_mode'],
noise_rate=opt['datasets']['val']['noise_rate'])
model.test()
visuals = model.get_current_visuals()
sr_img = util.tensor2img(visuals['rlt']) # uint8
gt_img = util.tensor2img(visuals['GT']) # 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
sr_img, gt_img = util.crop_border([sr_img, gt_img],
opt['scale'])
avg_psnr += util.calculate_psnr(sr_img, gt_img)
pbar.update('Test {}'.format(img_name))
avg_psnr = avg_psnr / idx
# log
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
if opt['use_wandb_logger'] and 'debug' not in opt['name']:
wandb.log({'psnr': avg_psnr}, step=current_step)
else: # video restoration validation
if opt['dist']:
# multi-GPU testing
psnr_rlt = {} # with border and center frames
if rank == 0:
pbar = util.ProgressBar(len(val_set))
for idx in range(rank, len(val_set), world_size):
val_data = val_set[idx]
val_data['LQs'].unsqueeze_(0)
val_data['GT'].unsqueeze_(0)
folder = val_data['folder']
idx_d, max_idx = val_data['idx'].split('/')
idx_d, max_idx = int(idx_d), int(max_idx)
if psnr_rlt.get(folder, None) is None:
psnr_rlt[folder] = torch.zeros(
max_idx,
dtype=torch.float32,
device='cuda')
# tmp = torch.zeros(max_idx, dtype=torch.float32, device='cuda')
model.feed_data(val_data,
noise_mode=opt['datasets']['val']
['noise_mode'],
noise_rate=opt['datasets']['val']
['noise_rate'])
model.test()
visuals = model.get_current_visuals()
rlt_img = util.tensor2img(visuals['rlt']) # uint8
gt_img = util.tensor2img(visuals['GT']) # uint8
# calculate PSNR
psnr_rlt[folder][idx_d] = util.calculate_psnr(
rlt_img, gt_img)
if rank == 0:
for _ in range(world_size):
pbar.update('Test {} - {}/{}'.format(
folder, idx_d, max_idx))
# # collect data
for _, v in psnr_rlt.items():
dist.reduce(v, 0)
dist.barrier()
if rank == 0:
psnr_rlt_avg = {}
psnr_total_avg = 0.
for k, v in psnr_rlt.items():
psnr_rlt_avg[k] = torch.mean(v).cpu().item()
psnr_total_avg += psnr_rlt_avg[k]
psnr_total_avg /= len(psnr_rlt)
log_s = '# Validation # PSNR: {:.4e}:'.format(
psnr_total_avg)
for k, v in psnr_rlt_avg.items():
log_s += ' {}: {:.4e}'.format(k, v)
logger.info(log_s)
if opt['use_tb_logger'] and 'debug' not in opt[
'name']:
tb_logger.add_scalar('psnr_avg',
psnr_total_avg,
current_step)
for k, v in psnr_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
if opt['use_wandb_logger'] and 'debug' not in opt[
'name']:
lq_img, rlt_img, gt_img = map(
util.tensor2img, [
visuals['LQ'], visuals['rlt'],
visuals['GT']
])
wandb.log({'psnr_avg': psnr_total_avg},
step=current_step)
wandb.log(psnr_rlt_avg, step=current_step)
wandb.log(
{
'Validation Image': [
wandb.Image(lq_img[:, :,
[2, 1, 0]],
caption='LQ'),
wandb.Image(rlt_img[:, :,
[2, 1, 0]],
caption='output'),
wandb.Image(gt_img[:, :,
[2, 1, 0]],
caption='GT'),
]
},
step=current_step)
else:
pbar = util.ProgressBar(len(val_loader))
psnr_rlt = {} # with border and center frames
psnr_rlt_avg = {}
psnr_total_avg = 0.
for val_data in val_loader:
folder = val_data['folder'][0]
idx_d = val_data['idx'].item()
# border = val_data['border'].item()
if psnr_rlt.get(folder, None) is None:
psnr_rlt[folder] = []
model.feed_data(val_data,
noise_mode=opt['datasets']['val']
['noise_mode'],
noise_rate=opt['datasets']['val']
['noise_rate'])
model.test()
visuals = model.get_current_visuals()
rlt_img = util.tensor2img(visuals['rlt']) # uint8
gt_img = util.tensor2img(visuals['GT']) # uint8
# calculate PSNR
psnr = util.calculate_psnr(rlt_img, gt_img)
psnr_rlt[folder].append(psnr)
pbar.update('Test {} - {}'.format(folder, idx_d))
for k, v in psnr_rlt.items():
psnr_rlt_avg[k] = sum(v) / len(v)
psnr_total_avg += psnr_rlt_avg[k]
psnr_total_avg /= len(psnr_rlt)
log_s = '# Validation # PSNR: {:.4e}:'.format(
psnr_total_avg)
for k, v in psnr_rlt_avg.items():
log_s += ' {}: {:.4e}'.format(k, v)
logger.info(log_s)
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr_avg', psnr_total_avg,
current_step)
for k, v in psnr_rlt_avg.items():
tb_logger.add_scalar(k, v, current_step)
if opt['use_wandb_logger'] and 'debug' not in opt[
'name']:
lq_img, rlt_img, gt_img = map(
util.tensor2img,
[visuals['LQ'], visuals['rlt'], visuals['GT']])
wandb.log({'psnr_avg': psnr_total_avg},
step=current_step)
wandb.log(psnr_rlt_avg, step=current_step)
wandb.log(
{
'Validation Image': [
wandb.Image(lq_img[:, :, [2, 1, 0]],
caption='LQ'),
wandb.Image(rlt_img[:, :, [2, 1, 0]],
caption='output'),
wandb.Image(gt_img[:, :, [2, 1, 0]],
caption='GT'),
]
},
step=current_step)
#### save models and training states
if current_step % opt['logger']['save_checkpoint_freq'] == 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.')
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.close()
#### mkdir and logger
util.mkdirs((path for key, path in opt_P['path'].items()
if not key == 'experiments_root' and 'pretrain_model' not in key
and 'resume' not in key))
util.mkdirs((path for key, path in opt_C['path'].items()
if not key == 'experiments_root' and 'pretrain_model' not in key
and 'resume' not in key))
util.setup_logger('base',
opt_P['path']['log'],
'test_' + opt_P['name'],
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt_P))
logger.info(option.dict2str(opt_C))
#### Create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt_P['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)
# load pretrained model by default
model_F = create_model(opt_F)
model_P = create_model(opt_P)
model_C = create_model(opt_C)
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:
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))
# 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
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
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)
#### 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.benchmark = 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 = create_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)
for _, train_data in enumerate(train_loader):
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)
#### log
if current_step % opt['logger']['print_freq'] == 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)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
# validation
if current_step % opt['train']['val_freq'] == 0 and rank <= 0:
avg_psnr = 0.0
idx = 0
for val_data in val_loader:
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
# 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
crop_size = opt['scale']
gt_img = gt_img / 255.
sr_img = sr_img / 255.
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.calculate_psnr(cropped_sr_img * 255,
cropped_gt_img * 255)
avg_psnr = avg_psnr / idx
# log
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
logger_val = logging.getLogger('val') # validation logger
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}'.format(
epoch, current_step, avg_psnr))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
#### save models and training states
if current_step % opt['logger']['save_checkpoint_freq'] == 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.')
def main(jsonPath):
# options
opt = option.parse(jsonPath, is_train=False)
util.mkdirs((path for key, path in opt["path"].items()
if not key == "pretrain_model_G"))
opt = option.dict_to_nonedict(opt)
util.setup_logger(None,
opt["path"]["log"],
"test.log",
level=logging.INFO,
screen=True)
logger = logging.getLogger("base")
logger.info(option.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)
logger.info("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"]
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
model.feed_data(data, need_HR=need_HR)
img_path = data["LR_path"][0]
img_name = os.path.splitext(os.path.basename(img_path))[0]
model.test() # test
visuals = model.get_current_visuals(need_HR=need_HR)
sr_img = util.tensor2img(visuals["SR"]) # 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)
# calculate PSNR and SSIM
if need_HR:
gt_img = util.tensor2img(visuals["HR"])
gt_img = gt_img / 255.0
sr_img = sr_img / 255.0
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)
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,
required=True,
help='Path to options file.')
opt = option.parse(parser.parse_args().opt, is_train=False)
util.mkdirs((path for key, path in opt['path'].items()
if not key == 'pretrain_model_G'))
opt = option.dict_to_nonedict(opt)
util.setup_logger(None,
opt['path']['log'],
'test.log',
level=logging.INFO,
screen=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt))
# Create test dataset and dataloader
test_loaders = []
znorm = 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['znorm'] and znorm == False:
znorm = True
# 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
model.feed_data(data, need_HR=need_HR)
img_path = data['LR_path'][0]
img_name = os.path.splitext(os.path.basename(img_path))[0]
model.test() # test
visuals = model.get_current_visuals(need_HR=need_HR)
if znorm: #opt['datasets']['train']['znorm']: # If the image range is [-1,1] # In testing, each "dataset" can have a different name (not train, val or other)
sr_img = util.tensor2img(visuals['SR'],
min_max=(-1, 1)) # uint8
else: # Default: Image range is [0,1]
sr_img = util.tensor2img(visuals['SR']) # 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)
# calculate PSNR and SSIM
if need_HR:
if znorm: #opt['datasets']['train']['znorm']: # If the image range is [-1,1] # In testing, each "dataset" can have a different name (not train, val or other)
gt_img = util.tensor2img(visuals['HR'],
min_max=(-1, 1)) # uint8
else: # Default: Image range is [0,1]
gt_img = util.tensor2img(visuals['HR']) # 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)
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, required=True, help='Path to options JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=False)
util.mkdirs((path for key, path in opt['path'].items() if not key == 'pretrain_model_G'))
opt = option.dict_to_nonedict(opt)
util.setup_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt))
# torch.backends.cudnn.benchmark = False
# torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = False
test_loaders = []
for phase, dataset_opt in opt['datasets'].items():
# dataset_opt['uids'] = uids[85:] if dataset_opt['data_type'] == 'h5' else uids
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set, dataset_opt)
logger.info('Number of test volumes in [{:s}]: {:d}'.format(dataset_opt['name'], len(test_set)))
test_loaders.append(test_loader)
# Create model
model = create_model(opt)
# copy net and convert to fp16 if necessary
model.half()
# create pdist model vgg
pdist_model = util.create_pdist_model(use_gpu=opt['gpu_ids'] is not None)
if opt["precision"] == 'int8':
model.prepare_quant(test_loaders[0])
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
logger.info('\nTesting [{:s}]...'.format(test_set_name))
results_dir = os.path.join(opt['path']['results_root'], test_set_name)
util.mkdir(results_dir)
pnsr_results = OrderedDict()
ssim_results = OrderedDict()
pdist_results = OrderedDict()
total_t = 0.
for i, data in enumerate(test_loader):
need_HR = False if test_loader.dataset.opt['dataroot_HR'] is None else True
has_mask = False if test_loader.dataset.opt['maskroot_HR'] is None else True
logger.info('start inference...')
model.feed_test_data(data, need_HR=need_HR)
t0 = time.time()
model.test(data) # test
if opt['gpu_ids']:
torch.cuda.synchronize()
t1 = time.time()
t = t1 - t0
total_t += t
logger.info('inference time: {:.3f}'.format(t))
# get cpu numpy from cuda tensor
visuals = model.get_current_visuals(data, maskOn=has_mask, need_HR=need_HR)
# save volume data
patient_id = data['uid'][0]
vol_path = os.path.join(results_dir, patient_id)
LR_spacings = [x.item() for x in data['spacings']]
if 'nrrd' in opt['result_format']:
logger.info('saving nnrd...')
sr_vol = util.tensor2img(visuals['SR'], out_type=np.uint16)
util.save_vol(opt, LR_spacings, vol_path + '.nrrd', sr_vol)
if 'dicom' in opt['result_format']:
logger.info('saving dicoms...')
sr_vol = util.tensor2img(visuals['SR'], out_type=np.int16, intercept = -1000)
util.save_dicoms(opt, LR_spacings, vol_path, sr_vol)
if opt['result_format']!='nrrd' and opt['result_format']!='dicom':
raise NotImplementedError('supported output format: nrrd or dicom')
#initialize dictionary with empty {}
pnsr_results[patient_id] = {}
ssim_results[patient_id] = {}
pdist_results[patient_id] = {}
# need_HR = False
if need_HR:
def _calculate_metrics(sr_vol, gt_vol, view='xy'):
sum_psnr = 0.
sum_ssim = 0.
sum_pdist = 0.
# [D,H,W]
num_val = 0 # psnr could be inf at xz or yz (near edges), will not calculate
for i, vol in enumerate(zip(sr_vol, gt_vol)):
sr_img, gt_img = vol[0], vol[1]
# calculate PSNR and SSIM
# range is assume to be [0,255] so have to scale back from 1500 to 255 float64
crop_size = round(opt['scale'])
cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size]\
.astype(np.float64) / 1500. * 255.
cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size]\
.astype(np.float64) / 1500. * 255.
psnr = util.calculate_psnr(cropped_sr_img, cropped_gt_img)
ssim = util.calculate_ssim(cropped_sr_img, cropped_gt_img)
if opt['datasets']['val']['get_pdist']:
pdist = util.calculate_pdist(pdist_model, cropped_sr_img, cropped_gt_img)
else:
pdist = float('nan')
if psnr != float('inf'):
num_val += 1
sum_psnr += psnr
sum_ssim += ssim
sum_pdist += pdist
logger.info('{:20s} - {:3d}- PSNR: {:.6f} dB; SSIM: {:.6f}; pdist: {:.6f}.'\
.format(patient_id, i+1, psnr, ssim, pdist))
pnsr_results[patient_id][view] = sum_psnr / num_val
ssim_results[patient_id][view] = sum_ssim / num_val
pdist_results[patient_id][view] = sum_pdist / num_val
return pnsr_results, ssim_results, pdist_results
sr_vol = util.tensor2img(visuals['SR'], out_type=np.uint16)
gt_vol = util.tensor2img(visuals['HR'], out_type=np.uint16) # uint16 range [0,1500]
min_depth = min(sr_vol.shape[0], gt_vol.shape[0])
sr_vol = sr_vol[:min_depth,...]
gt_vol = gt_vol[:min_depth,...] # make sure they have the same depth
# [H W] axial view
_calculate_metrics(sr_vol, gt_vol, view='xy')
# [D W] coronal view
_calculate_metrics(sr_vol.transpose(1, 0, 2), gt_vol.transpose(1, 0, 2), view='xz')
# [D H] sagittal view
_calculate_metrics(sr_vol.transpose(2, 0, 1), gt_vol.transpose(2, 0, 1), view='yz')
else:
logger.info(patient_id)
if need_HR: # metrics
# print result dictionary
util.print_metrics(logger, 'test PSNR', pnsr_results)
util.print_metrics(logger, 'test SSIM', ssim_results)
util.print_metrics(logger, 'test pdist', pdist_results)
logger.info('average inference time: {:.3f}'.format(total_t/len(test_loader)))
parser.add_argument('--opt',
type=str,
default="options/test/test_ESRGAN.json",
help='Path to options JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=False)
util.mkdirs((path for key, path in opt['path'].items()
if not key == 'pretrain_model_G'))
opt = option.dict_to_nonedict(opt)
util.setup_logger(None,
opt['path']['log'],
'test.log',
level=logging.INFO,
screen=True)
logger = logging.getLogger('base')
logger.info(option.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)
logger.info('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']
logger.info('\nTesting [{:s}]...'.format(test_set_name))
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option JSON file.')
parser.add_argument('-name', type=str, help='Name of the experiment.')
parser.add_argument('-seed', type=int, default=2, help='Random seed.')
parser.add_argument('-gsize',
type=int,
default=1,
help='Num of devices per task.')
parser.add_argument('-val_freq',
type=int,
help='Num of transfer frequency.')
parser.add_argument('-train',
action='store_true',
help='Specify if training.')
parser.add_argument(
'-lmdb',
action='store_true',
help='create or use lmdb datasets for accelerating I/O.')
parser.add_argument('-droot', type=str, help='dataroot if need specify.')
parser.add_argument('-rancl',
action='store_true',
help='random task for transfer?')
parser.add_argument('-create_val',
action='store_true',
help='create validation set from training set')
parser.add_argument('-model', type=str, help='specify model to use')
parser.add_argument('-k',
action='store_true',
help='Use knowledge distillation based transfer.')
parser.add_argument('-a',
action='store_true',
help='Use attention based transfer.')
parser.add_argument('-f',
action='store_true',
help='Use fsp matrix based transfer.')
parser.add_argument('-w',
action='store_true',
help='Use weights transfer.')
parser.add_argument('-ws',
action='store_true',
help='Use weights statistical transfer.')
parser.add_argument('-VL',
action='store_true',
help='Variate on loss function.')
parser.add_argument('-VO',
action='store_true',
help='Variate on optimization algorithm.')
parser.add_argument('-VH',
action='store_true',
help='Variate on hyperparameters.')
parser.add_argument('-VD',
action='store_true',
help='Variate on datasets.')
parser.add_argument('-VF',
action='store_true',
help='Variate on cross validation.')
parser.add_argument('-VS', type=float, help='fraction of bagging dataset.')
opt = parser.parse_args()
if not opt.train and opt.gsize > 1:
warnings.warn(
'1 device for validation is enough! The gsize will be reset to 1')
opt.gsize = 1
if not opt.w and opt.gsize > 1:
warnings.warn(
'The exploit-explorer mode is not activated! The gsize will be reset to 1'
)
opt.gsize = 1
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
gpus_per_node = torch.cuda.device_count()
if gpus_per_node == 0:
device_id = 'cpu'
else:
device_id = rank % gpus_per_node
print('start running! {} gpus are in use'.format(gpus_per_node))
print('Torch version {}'.format(torch.__version__))
# run to here is fine!!
# torch.cuda.set_device(device_id)
device_per_task = opt.gsize
task_id = rank // device_per_task
task_exploit_rank = rank - rank % device_per_task
world_size = comm.Get_size()
opt.ntasks = world_size
opt = option.parse(opt, task_id)
opt = option.dict_to_nonedict(
opt) # Convert to NoneDict, which return None for missing key.
# Setup directory
if not opt['resume'] and task_exploit_rank == rank and not opt['path'][
'resume_state']:
try:
if opt['is_train'] and rank == 0:
util.mkdir_and_rename(
opt['path']
['experiments_root']) # rename old folder if exists
except FileNotFoundError:
raise FileNotFoundError("Issue from task {} and rank {}".format(
task_id, rank))
comm.Barrier()
if not opt['resume'] and task_exploit_rank == rank and not opt['path'][
'resume_state']:
util.mkdirs((path for key, path in opt['path'].items() if key not in [
'experiments_root', 'log', 'root', 'pretrain_model',
'pretrain_model_G', 'resume_state', 'data_config'
]))
# save json file to task folder
if opt['is_train']:
with open(os.path.join(opt['path']['task'], 'task.json'),
'w') as outfile:
json.dump(opt, outfile, indent=4, ensure_ascii=False)
comm.Barrier()
# config loggers. Before it, the log will not work
util.setup_logger(str(task_id),
opt['path']['log'],
'train',
level=logging.INFO,
screen=True,
task_id=task_id,
rank=rank)
logger = logging.getLogger(str(task_id))
if task_exploit_rank == rank:
if task_id == 0:
logger.info(option.dict2str(opt)) # display options
else:
logger.info(
'Auxiliary task {} configuration: network: {}, optim: {}, loss: {}, data: {}'
.format(task_id, opt['network'], opt['train']['optim'],
opt['train']['loss'],
opt['datasets']['train']['name']))
# tensorboard logger
tb_logger = None
if opt['use_tb_logger'] and opt[
'is_train'] and task_exploit_rank == rank: # and 'debug' not in opt['name']
from tensorboardX import SummaryWriter
if rank == 0:
util.mkdir_and_rename(os.path.join('../tb_logger/', opt['name']))
comm.barrier()
tb_logger_path = os.path.join('../tb_logger/', opt['name'],
str(task_id))
tb_logger = SummaryWriter(log_dir=tb_logger_path)
# create task
# If you don't use the following setting, the results will still be non-deterministic.
# However, if you set them so, the speed may be slower, depending on models.
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
task = create_task(opt, comm=comm, device=device_id)
# Start training or testing
start = time()
if opt['is_train']:
task.resume_training()
task.full_training(opt['logger']['print_freq'], opt['val_freq'],
opt['logger']['save_checkpoint_freq'], tb_logger)
if rank == 0:
if opt['varyOnCV']:
all_hmean = task.cv_validation()
results = [r['acc'] for r in all_hmean]
idx = np.argmax(results)
best_task_id = all_hmean[idx]['id']
best_task_rank = device_per_task * best_task_id
logger.info(
'All hmean is {} and select rank {} to run validation'.
format(all_hmean, best_task_rank))
if task_exploit_rank == best_task_rank:
task.validation(verbose=True,
report=True,
split='test',
best=True)
else:
task.validation(verbose=True,
report=True,
split='test',
best=True)
logger.info('End of training.')
else:
logger.info('Task {} start validation'.format(task_id))
task.validation(verbose=True, report=True, split='test', best=True)
duration = time() - start
logger.info('The program tasks time {}'.format(duration))
