import argparse
import json
import os
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import yaml
import torch
import torch.nn as nn
from tqdm import tqdm
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import MultiStepLR, CosineAnnealingLR
import datasets
import models
import utils
def make_data_loader(spec, tag='', local_rank=0):
if spec is None:
return None
dataset = datasets.make(spec['dataset'])
dataset = datasets.make(spec['wrapper'], args={'dataset': dataset})
if local_rank == 0:
print('{} dataset: size={}'.format(tag, len(dataset)))
for k, v in dataset[0].items():
if torch.is_tensor(v):
print(' {}: shape={}'.format(k, v.shape))
elif isinstance(v, str):
pass
elif isinstance(v, dict):
for k0, v0 in v.items():
if hasattr(v0, 'shape'):
print(' {}: shape={}'.format(k0, v0.shape))
else:
raise NotImplementedError
sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=(tag == 'train'))
loader = torch.utils.data.DataLoader(dataset,
batch_size=spec['batch_size'],
num_workers=spec['num_workers'],
pin_memory=True,
sampler=sampler)
return loader
def make_data_loaders(config, local_rank):
train_loader = make_data_loader(config.get('train_dataset'), tag='train', local_rank=local_rank)
val_loader = make_data_loader(config.get('val_dataset'), tag='val', local_rank=local_rank)
return train_loader, val_loader
def prepare_training(config, local_rank):
if config.get('resume') is not None:
sv_file = torch.load(config['resume'])
model = models.make(sv_file['model'], load_sd=True).cuda()
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
optimizer = utils.make_optimizer(
model.parameters(), sv_file['optimizer'], load_sd=True)
epoch_start = sv_file['epoch'] + 1
if config.get('multi_step_lr') is None:
lr_scheduler = None
else:
lr_scheduler = MultiStepLR(optimizer, **config['multi_step_lr'])
for _ in range(epoch_start - 1):
lr_scheduler.step()
else:
model = models.make(config['model']).cuda(local_rank)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
optimizer = utils.make_optimizer(
model.parameters(), config['optimizer'])
epoch_start = 1
lr_scheduler = config.get('lr_scheduler')
lr_scheduler_name = lr_scheduler.pop('name')
if 'MultiStepLR' == lr_scheduler_name:
lr_scheduler = MultiStepLR(optimizer, **lr_scheduler)
elif 'CosineAnnealingLR' == lr_scheduler_name:
lr_scheduler = CosineAnnealingLR(optimizer, **lr_scheduler)
elif 'CosineAnnealingWarmUpLR' == lr_scheduler_name:
lr_scheduler = utils.warm_up_cosine_lr_scheduler(optimizer, **lr_scheduler)
if local_rank == 0:
print('model: #params={}'.format(utils.compute_num_params(model, text=True)))
return model, optimizer, epoch_start, lr_scheduler
def reduce_mean(tensor, nprocs):
rt = tensor.clone()
dist.all_reduce(rt, op=dist.ReduceOp.SUM)
rt /= nprocs
return rt
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self, name, fmt=':f'):
self.name = name
self.fmt = fmt
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def return_avg(self):
return self.avg
def __str__(self):
fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
return fmtstr.format(**self.__dict__)
def train(train_loader, model, optimizer, local_rank):
model = model.train()
loss_fn = nn.L1Loss().cuda(local_rank)
train_losses = AverageMeter('Loss', ':.4e')
data_norm = config['data_norm']
t = data_norm['img']
img_sub = torch.FloatTensor(t['sub']).view(1, -1, 1, 1).cuda(local_rank)
img_div = torch.FloatTensor(t['div']).view(1, -1, 1, 1).cuda(local_rank)
t = data_norm['gt']
gt_sub = torch.FloatTensor(t['sub']).view(1, 1, -1).cuda(local_rank)
gt_div = torch.FloatTensor(t['div']).view(1, 1, -1).cuda(local_rank)
if local_rank == 0:
pbar = tqdm(total=len(train_loader), desc='train', leave=False)
for i, batch in enumerate(train_loader):
if local_rank == 0:
pbar.update(1)
keys = list(batch.keys())
batch = batch[keys[torch.randint(0, len(keys), [])]]
for k, v in batch.items():
if torch.is_tensor(v):
batch[k] = v.cuda(local_rank, non_blocking=True)
img = (batch['img'] - img_sub) / img_div
gt = (batch['gt'] - gt_sub) / gt_div
pred = model(img, gt.shape[-2:])
if isinstance(pred, tuple):
loss = 0.2 * loss_fn(pred[0], gt) + loss_fn(pred[1], gt)
elif isinstance(pred, list):
losses = [loss_fn(x, gt) for x in pred]
losses = [x * (idx + 1) for idx, x in enumerate(losses)]
loss = sum(losses) / ((1 + len(losses)) * len(losses) / 2)
else:
loss = loss_fn(pred, gt)
torch.distributed.barrier()
reduced_loss = reduce_mean(loss, dist.get_world_size())
train_losses.update(reduced_loss.item(), img.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
if local_rank == 0:
pbar.close()
return train_losses.avg
def eval_psnr(loader, class_names, model, local_rank, data_norm=None, eval_type=None, eval_bsize=None, verbose=False, crop_border=4):
crop_border = int(crop_border) if crop_border else crop_border
if local_rank == 0:
print('crop border: ', crop_border)
model = model.eval()
if data_norm is None:
data_norm = {
'img': {'sub': [0], 'div': [1]},
'gt': {'sub': [0], 'div': [1]}
}
t = data_norm['img']
img_sub = torch.FloatTensor(t['sub']).view(1, -1, 1, 1).cuda(local_rank)
img_div = torch.FloatTensor(t['div']).view(1, -1, 1, 1).cuda(local_rank)
t = data_norm['gt']
gt_sub = torch.FloatTensor(t['sub']).view(1, -1, 1, 1).cuda(local_rank)
gt_div = torch.FloatTensor(t['div']).view(1, -1, 1, 1).cuda(local_rank)
if eval_type is None:
metric_fn = [utils.calculate_psnr_pt, utils.calculate_ssim_pt]
elif eval_type == 'psnr+ssim':
metric_fn = [utils.calculate_psnr_pt, utils.calculate_ssim_pt]
elif eval_type.startswith('div2k'):
scale = int(eval_type.split('-')[1])
metric_fn = partial(utils.calc_psnr, dataset='div2k', scale=scale)
elif eval_type.startswith('benchmark'):
scale = int(eval_type.split('-')[1])
metric_fn = partial(utils.calc_psnr, dataset='benchmark', scale=scale)
else:
raise NotImplementedError
val_res_psnr = AverageMeter('psnr', ':.4f')
val_res_ssim = AverageMeter('ssim', ':.4f')
if local_rank == 0:
pbar = tqdm(total=len(loader), desc='val', leave=False)
for batch in loader:
if local_rank == 0:
pbar.update(1)
for k, v in batch.items():
if torch.is_tensor(v):
batch[k] = v.cuda(local_rank, non_blocking=True)
img = (batch['img'] - img_sub) / img_div
with torch.no_grad():
pred = model(img, batch['gt'].shape[-2:])
if isinstance(pred, list):
pred = pred[-1]
pred = pred * gt_div + gt_sub
res_psnr = metric_fn[0](
pred,
batch['gt'],
crop_border=crop_border
).mean()
res_ssim = metric_fn[1](
pred,
batch['gt'],
crop_border=crop_border
).mean()
torch.distributed.barrier()
reduced_val_res_psnr = reduce_mean(res_psnr, dist.get_world_size())
reduced_val_res_ssim = reduce_mean(res_ssim, dist.get_world_size())
val_res_psnr.update(reduced_val_res_psnr.item(), img.size(0))
val_res_ssim.update(reduced_val_res_ssim.item(), img.size(0))
if verbose and local_rank == 0:
pbar.set_description(
'val psnr: {:.4f} ssim: {:.4f}'.format(val_res_psnr.avg, val_res_ssim.avg))
if local_rank == 0:
pbar.close()
return val_res_psnr.avg, val_res_ssim.avg
def main(config, save_path):
# torch.backends.cudnn.benchmark = True
dist.init_process_group("nccl")
rank = dist.get_rank()
local_rank = int(os.environ["LOCAL_RANK"])
world_size = dist.get_world_size()
print(f'rank: {rank} local_rank: {local_rank} world_size: {world_size}')
# print(f'local_rank: {torch.distributed.local_rank()}')
if local_rank == 0:
log, writer = utils.set_save_path(save_path)
with open(os.path.join(save_path, 'config.yaml'), 'w') as f:
yaml.dump(config, f, sort_keys=False)
train_loader, val_loader = make_data_loaders(config, local_rank)
if config.get('data_norm') is None:
config['data_norm'] = {
'img': {'sub': [0], 'div': [1]},
'gt': {'sub': [0], 'div': [1]}
}
model, optimizer, epoch_start, lr_scheduler = prepare_training(config, local_rank)
epoch_max = config['epoch_max']
epoch_val_interval = config.get('epoch_val_interval')
epoch_save_interval = config.get('epoch_save_interval')
max_val_v = -1e18
timer = utils.Timer()
for epoch in range(epoch_start, epoch_max + 1):
t_epoch_start = timer.t()
train_loader.sampler.set_epoch(epoch)
train_loss = train(train_loader, model, optimizer, local_rank)
if lr_scheduler is not None:
lr_scheduler.step()
if rank == 0:
log_info = ['epoch {}/{}'.format(epoch, epoch_max)]
log_info.append('train: loss={:.4f}'.format(train_loss))
writer.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch)
writer.add_scalars('loss', {'train': train_loss}, epoch)
model_ = model.module
model_spec = config['model']
model_spec['sd'] = model_.state_dict()
optimizer_spec = config['optimizer']
optimizer_spec['sd'] = optimizer.state_dict()
sv_file = {
'model': model_spec,
'optimizer': optimizer_spec,
'epoch': epoch
}
if rank == 0:
torch.save(sv_file, os.path.join(save_path, 'epoch-last.pth'))
if (epoch_save_interval is not None) and (epoch % epoch_save_interval == 0):
if rank == 0:
torch.save(sv_file, os.path.join(save_path, 'epoch-{}.pth'.format(epoch)))
if (epoch_val_interval is not None) and (epoch % epoch_val_interval == 0):
file_names = json.load(open(config['val_dataset']['dataset']['args']['split_file']))['test']
class_names = list(set([os.path.basename(os.path.dirname(x)) for x in file_names]))
val_res_psnr, val_res_ssim = eval_psnr(val_loader, class_names, model_, local_rank,
data_norm=config['data_norm'],
eval_type=config.get('eval_type'),
eval_bsize=config.get('eval_bsize'),
crop_border=4)
if rank == 0:
log_info.append('val: psnr={:.4f}'.format(val_res_psnr))
writer.add_scalars('psnr', {'val': val_res_psnr}, epoch)
if val_res_psnr > max_val_v:
max_val_v = val_res_psnr
if rank == 0:
torch.save(sv_file, os.path.join(save_path, 'epoch-best.pth'))
t = timer.t()
if rank == 0:
prog = (epoch - epoch_start + 1) / (epoch_max - epoch_start + 1)
t_epoch = utils.time_text(t - t_epoch_start)
t_elapsed, t_all = utils.time_text(t), utils.time_text(t / prog)
log_info.append('{} {}/{}'.format(t_epoch, t_elapsed, t_all))
log(', '.join(log_info))
writer.flush()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--config', default='configs/train_1x-5x_INR_funsr.yaml')
parser.add_argument('--name', default='EXP20221216_11')
parser.add_argument('--tag', default=None)
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
print('config loaded.')
save_name = args.name
if save_name is None:
save_name = '_' + args.config.split('/')[-1][:-len('.yaml')]
if args.tag is not None:
save_name += '_' + args.tag
save_path = os.path.join('./checkpoints', save_name)
main(config, save_path)