init + inference.py патч
This commit is contained in:
221
train.py
Normal file
221
train.py
Normal file
@ -0,0 +1,221 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user