221 lines
10 KiB
Python
221 lines
10 KiB
Python
|
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)
|