import os
import cv2
import time
import random
import datetime
import argparse
import numpy as np
from tqdm import tqdm
from piq import ssim,psnr
from itertools import cycle
import torch
import torch.nn as nn
from torch.utils import data
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from utils import dict2string,mkdir,get_lr,torch2cvimg,second2hours
from loaders import docres_loader
from models import restormer_arch
def seed_torch(seed=1029):
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
#torch.use_deterministic_algorithms(True)
# seed_torch()
def getBasecoord(h,w):
base_coord0 = np.tile(np.arange(h).reshape(h,1),(1,w)).astype(np.float32)
base_coord1 = np.tile(np.arange(w).reshape(1,w),(h,1)).astype(np.float32)
base_coord = np.concatenate((np.expand_dims(base_coord1,-1),np.expand_dims(base_coord0,-1)),-1)
return base_coord
def train(args):
## DDP init
dist.init_process_group(backend='nccl',init_method='env://',timeout=datetime.timedelta(seconds=36000))
torch.cuda.set_device(args.local_rank)
device = torch.device('cuda',args.local_rank)
torch.cuda.manual_seed_all(42)
### Log file:
mkdir(args.logdir)
mkdir(os.path.join(args.logdir,args.experiment_name))
log_file_path=os.path.join(args.logdir,args.experiment_name,'log.txt')
log_file=open(log_file_path,'a')
log_file.write('\n--------------- '+args.experiment_name+' ---------------\n')
log_file.close()
### Setup tensorboard for visualization
if args.tboard:
writer = SummaryWriter(os.path.join(args.logdir,args.experiment_name,'runs'),args.experiment_name)
### Setup Dataloader
datasets_setting = [
{'task':'deblurring','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/deblurring/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/deblurring/tdd/train.json']},
{'task':'dewarping','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/dewarping/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/dewarping/doc3d/train_1_19.json']},
{'task':'binarization','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/binarization/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/binarization/train.json']},
{'task':'deshadowing','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/deshadowing/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/deshadowing/train.json']},
{'task':'appearance','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/appearance/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/appearance/trainv2.json']}
]
ratios = [dataset_setting['ratio'] for dataset_setting in datasets_setting]
datasets = [docres_loader.DocResTrainDataset(dataset=dataset_setting,img_size=args.im_size) for dataset_setting in datasets_setting]
trainloaders = [{'task':datasets_setting[i],'loader':data.DataLoader(dataset=datasets[i], sampler=DistributedSampler(datasets[i]), batch_size=args.batch_size, num_workers=2, pin_memory=True,drop_last=True),'iter_loader':iter(data.DataLoader(dataset=datasets[i], sampler=DistributedSampler(datasets[i]), batch_size=args.batch_size, num_workers=2, pin_memory=True,drop_last=True))} for i in range(len(datasets))]
### test loader
# for i in tqdm(range(args.total_iter)):
# loader_index = random.choices(list(range(len(trainloaders))),ratios)[0]
# in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])
### Setup Model
model = restormer_arch.Restormer(
inp_channels=6,
out_channels=3,
dim = 48,
num_blocks = [2,3,3,4],
num_refinement_blocks = 4,
heads = [1,2,4,8],
ffn_expansion_factor = 2.66,
bias = False,
LayerNorm_type = 'WithBias',
dual_pixel_task = True
)
model=DDP(model.cuda(),device_ids=[args.local_rank],output_device=args.local_rank)
### Optimizer
optimizer= torch.optim.AdamW(model.parameters(),lr=args.l_rate,weight_decay=5e-4)
### LR Scheduler
sched = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.total_iter, eta_min=1e-6, last_epoch=-1)
### load checkpoint
iter_start=0
if args.resume is not None:
print("Loading model and optimizer from checkpoint '{}'".format(args.resume))
x = checkpoint['model_state']
model.load_state_dict(x,strict=False)
iter_start=checkpoint['iter']
print("Loaded checkpoint '{}' (iter {})".format(args.resume, iter_start))
###-----------------------------------------Training-----------------------------------------
##initialize
scaler = torch.cuda.amp.GradScaler()
loss_dict = {}
total_step = 0
l2 = nn.MSELoss()
l1 = nn.L1Loss()
ce = nn.CrossEntropyLoss()
bce = nn.BCEWithLogitsLoss()
m = nn.Sigmoid()
best = 0
best_ce = 999
## total_steps
for iters in range(iter_start,args.total_iter):
start_time = time.time()
loader_index = random.choices(list(range(len(trainloaders))),ratios)[0]
try:
in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])
except StopIteration:
trainloaders[loader_index]['iter_loader']=iter(trainloaders[loader_index]['loader'])
in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])
in_im = in_im.float().cuda()
gt_im = gt_im.float().cuda()
binarization_loss,appearance_loss,dewarping_loss,deblurring_loss,deshadowing_loss = 0,0,0,0,0
with torch.cuda.amp.autocast():
pred_im = model(in_im,trainloaders[loader_index]['task']['task'])
if trainloaders[loader_index]['task']['task'] == 'binarization':
gt_im = gt_im.long()
binarization_loss = ce(pred_im[:,:2,:,:], gt_im[:,0,:,:])
loss = binarization_loss
elif trainloaders[loader_index]['task']['task'] == 'dewarping':
dewarping_loss = l1(pred_im[:,:2,:,:], gt_im[:,:2,:,:])
loss = dewarping_loss
elif trainloaders[loader_index]['task']['task'] == 'appearance':
appearance_loss = l1(pred_im, gt_im)
loss = appearance_loss
elif trainloaders[loader_index]['task']['task'] == 'deblurring':
deblurring_loss = l1(pred_im, gt_im)
loss = deblurring_loss
elif trainloaders[loader_index]['task']['task'] == 'deshadowing':
deshadowing_loss = l1(pred_im, gt_im)
loss = deshadowing_loss
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
loss_dict['dew_loss']=dewarping_loss.item() if isinstance(dewarping_loss,torch.Tensor) else 0
loss_dict['app_loss']=appearance_loss.item() if isinstance(appearance_loss,torch.Tensor) else 0
loss_dict['des_loss']=deshadowing_loss.item() if isinstance(deshadowing_loss,torch.Tensor) else 0
loss_dict['deb_loss']=deblurring_loss.item() if isinstance(deblurring_loss,torch.Tensor) else 0
loss_dict['bin_loss']=binarization_loss.item() if isinstance(binarization_loss,torch.Tensor) else 0
end_time = time.time()
duration = end_time-start_time
## log
if (iters+1) % 10 == 0:
## print
print('iters [{}/{}] -- '.format(iters+1,args.total_iter)+dict2string(loss_dict)+' --lr {:6f}'.format(get_lr(optimizer))+' -- time {}'.format(second2hours(duration*(args.total_iter-iters))))
## tbord
if args.tboard:
for key,value in loss_dict.items():
writer.add_scalar('Train '+key+'/Iterations', value, total_step)
## logfile
with open(log_file_path,'a') as f:
f.write('iters [{}/{}] -- '.format(iters+1,args.total_iter)+dict2string(loss_dict)+' --lr {:6f}'.format(get_lr(optimizer))+' -- time {}'.format(second2hours(duration*(args.total_iter-iters)))+'\n')
if (iters+1) % 5000 == 0:
state = {'iters': iters+1,
'model_state': model.state_dict(),
'optimizer_state' : optimizer.state_dict(),}
if not os.path.exists(os.path.join(args.logdir,args.experiment_name)):
os.system('mkdir ' + os.path.join(args.logdir,args.experiment_name))
if torch.distributed.get_rank()==0:
torch.save(state, os.path.join(args.logdir,args.experiment_name,"{}.pkl".format(iters+1)))
sched.step()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Hyperparams')
parser.add_argument('--im_size', nargs='?', type=int, default=256,
help='Height of the input image')
parser.add_argument('--total_iter', nargs='?', type=int, default=100000,
help='# of the epochs')
parser.add_argument('--batch_size', nargs='?', type=int, default=10,
help='Batch Size')
parser.add_argument('--l_rate', nargs='?', type=float, default=2e-4,
help='Learning Rate')
parser.add_argument('--resume', nargs='?', type=str, default=None,
help='Path to previous saved model to restart from')
parser.add_argument('--logdir', nargs='?', type=str, default='./checkpoints/',
help='Path to store the loss logs')
parser.add_argument('--tboard', dest='tboard', action='store_true',
help='Enable visualization(s) on tensorboard | False by default')
parser.add_argument('--local_rank',type=int,default=0,metavar='N')
parser.add_argument('--experiment_name', nargs='?', type=str,default='experiment_name',
help='the name of this experiment')
parser.set_defaults(tboard=False)
args = parser.parse_args()
train(args)