import os import datetime import argparse import numpy as np import torch import torch.optim as optim from torch.autograd import Variable from torch.utils.data import DataLoader from ImagePairPrefixFolder import ImagePairPrefixFolder, var_custom_collate from utils import MovingAvg from tf_visualizer import TFVisualizer parser = argparse.ArgumentParser() parser.add_argument('--network', default='GCANet') parser.add_argument('--name', default='default_exp') parser.add_argument('--gpu_ids', default='0') parser.add_argument('--epochs', type=int, default=100) parser.add_argument('--lr', type=float, default=0.001) parser.add_argument('--lr_step', type=int, default=40) parser.add_argument('--lr_gamma', type=float, default=0.1) parser.add_argument('--weight_decay', type=float, default=0.0005) parser.add_argument('--checkpoints_dir', default='checkpoint') parser.add_argument('--logDir', default='tblogdir') parser.add_argument('--resume_dir', default='') parser.add_argument('--resume_epoch', type=int, default=0) parser.add_argument('--save_epoch', type=int, default=5) parser.add_argument('--save_latest_freq', type=int, default=5000) parser.add_argument('--test_epoch', type=int, default=5) parser.add_argument('--test_max_size', type=int, default=1080) parser.add_argument('--size_unit', type=int, default=8) parser.add_argument('--print_iter', type=int, default=100) parser.add_argument('--input_folder', default='') parser.add_argument('--gt_folder', default='') parser.add_argument('--test_input_folder', default='') parser.add_argument('--test_gt_folder', default='') parser.add_argument('--num_workers', type=int, default=16) parser.add_argument('--batch_size', type=int, default=4) parser.add_argument('--only_residual', action='store_true', help='regress residual rather than image') parser.add_argument('--loss_func', default='l2', help='l2|l1') parser.add_argument('--inc', type=int, default=3) parser.add_argument('--outc', type=int, default=3) parser.add_argument('--force_rgb', action='store_true') parser.add_argument('--no_edge', action='store_true') opt = parser.parse_args() opt.input_folder = os.path.expanduser(opt.input_folder) opt.gt_folder = os.path.expanduser(opt.gt_folder) opt.test_input_folder = os.path.expanduser(opt.test_input_folder) opt.test_gt_folder = os.path.expanduser(opt.test_gt_folder) if not os.path.exists(os.path.join(opt.checkpoints_dir, opt.name)): os.makedirs(os.path.join(opt.checkpoints_dir, opt.name)) opt.resume_dir = opt.resume_dir if opt.resume_dir != '' else os.path.join(opt.checkpoints_dir, opt.name) visualizer = TFVisualizer(opt) ### Log out with open(os.path.realpath(__file__), 'r') as fid: visualizer.print_logs(fid.read()) ## print argument for key, val in vars(opt).items(): visualizer.print_logs('%s: %s' % (key, val)) opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(',')] assert all(0 <= x <= torch.cuda.device_count() for x in opt.gpu_ids), 'gpu id should ' \ 'be 0~{0}'.format(torch.cuda.device_count()) torch.cuda.set_device(opt.gpu_ids[0]) train_dataset = ImagePairPrefixFolder(opt.input_folder, opt.gt_folder, size_unit=opt.size_unit, force_rgb=opt.force_rgb) train_dataloader = DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, collate_fn=var_custom_collate, pin_memory=True, num_workers=opt.num_workers) opt.do_test = opt.test_gt_folder != '' if opt.do_test: test_dataset = ImagePairPrefixFolder(opt.test_input_folder, opt.test_gt_folder, max_img_size=opt.test_max_size, size_unit=opt.size_unit, force_rgb=opt.force_rgb) test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False, collate_fn=var_custom_collate, pin_memory=True, num_workers=1) total_inc = opt.inc if opt.no_edge else opt.inc + 1 if opt.network == 'GCANet': from GCANet import GCANet net = GCANet(in_c=total_inc, out_c=3, only_residual=opt.only_residual) else: print('network structure %s not supported' % opt.network) raise ValueError if opt.loss_func == 'l2': loss_crit = torch.nn.MSELoss() elif opt.loss_func == 'l1': loss_crit = torch.nn.SmoothL1Loss() else: print('loss_func %s not supported' % opt.loss_func) raise ValueError pnsr_crit = torch.nn.MSELoss() if len(opt.gpu_ids) > 0: net.cuda() if len(opt.gpu_ids) > 1: net = torch.nn.DataParallel(net) loss_crit = loss_crit.cuda() pnsr_crit = pnsr_crit.cuda() optimizer = optim.Adam(net.parameters(), lr=opt.lr) step_optim_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=opt.lr_step, gamma=opt.lr_gamma) loss_avg = MovingAvg(pool_size=50) start_epoch = 0 total_iter = 0 if os.path.exists(os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth')): print('resuming from latest.pth') latest_info = torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth')) start_epoch = latest_info['epoch'] total_iter = latest_info['total_iter'] if isinstance(net, torch.nn.DataParallel): net.module.load_state_dict(latest_info['net_state']) else: net.load_state_dict(latest_info['net_state']) optimizer.load_state_dict(latest_info['optim_state']) if opt.resume_epoch > 0: start_epoch = opt.resume_epoch total_iter = opt.resume_epoch * len(train_dataloader) resume_path = os.path.join(opt.resume_epoch, 'net_epoch_%d.pth') % opt.resume_epoch print('resume from : %s' % resume_path) assert os.path.exists(resume_path), 'cannot find the resume model: %s ' % resume_path if isinstance(net, torch.nn.DataParallel): net.module.load_state_dict(torch.load(resume_path)) else: net.load_state_dict(torch.load(resume_path)) for epoch in range(start_epoch, opt.epochs): visualizer.print_logs("Start to train epoch %d" % epoch) net.train() for iter, data in enumerate(train_dataloader): total_iter += 1 optimizer.zero_grad() step_optim_scheduler.step(epoch) batch_input_img, batch_input_edge, batch_gt = data if len(opt.gpu_ids) > 0: batch_input_img, batch_input_edge, batch_gt = batch_input_img.cuda(), batch_input_edge.cuda(), batch_gt.cuda() if opt.no_edge: batch_input = batch_input_img else: batch_input = torch.cat((batch_input_img, batch_input_edge), dim=1) batch_input_v = Variable(batch_input) if opt.only_residual: batch_gt_v = Variable(batch_gt - (batch_input_img+128)) else: batch_gt_v = Variable(batch_gt) pred = net(batch_input_v) loss = loss_crit(pred, batch_gt_v) avg_loss = loss_avg.set_curr_val(loss.data) loss.backward() optimizer.step() if iter % opt.print_iter == 0: visualizer.plot_current_losses(total_iter, { 'loss': loss}) visualizer.print_logs('%s Step[%d/%d], lr: %f, mv_avg_loss: %f, loss: %f' % (str(datetime.datetime.now()).split(' ')[1], iter, len(train_dataloader), step_optim_scheduler.get_lr()[0], avg_loss, loss)) if total_iter % opt.save_latest_freq == 0: latest_info = {'total_iter': total_iter, 'epoch': epoch, 'optim_state': optimizer.state_dict()} if len(opt.gpu_ids) > 1: latest_info['net_state'] = net.module.state_dict() else: latest_info['net_state'] = net.state_dict() print('save lastest model.') torch.save(latest_info, os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth')) if (epoch+1) % opt.save_epoch == 0 : visualizer.print_logs('saving model for epoch %d' % epoch) if len(opt.gpu_ids) > 1: torch.save(net.module.state_dict(), os.path.join(opt.checkpoints_dir, opt.name, 'net_epoch_%d.pth' % (epoch+1))) else: torch.save(net.state_dict(), os.path.join(opt.checkpoints_dir, opt.name, 'net_epoch_%d.pth' % (epoch + 1))) if opt.do_test: avg_psnr = 0 task_cnt = 0 net.eval() with torch.no_grad(): for iter, data in enumerate(test_dataloader): batch_input_img, batch_input_edge, batch_gt = data if len(opt.gpu_ids) > 0: batch_input_img, batch_input_edge, batch_gt = batch_input_img.cuda(), batch_input_edge.cuda(), batch_gt.cuda() if opt.no_edge: batch_input = batch_input_img else: batch_input = torch.cat((batch_input_img, batch_input_edge), dim=1) batch_input_v = Variable(batch_input) batch_gt_v = Variable(batch_gt) pred = net(batch_input_v) if opt.only_residual: loss = pnsr_crit(pred+Variable(batch_input_img+128), batch_gt_v) else: loss = pnsr_crit(pred, batch_gt_v) avg_psnr += 10 * np.log10(255 * 255 / loss.item()) task_cnt += 1 visualizer.print_logs('Testing for epoch: %d' % epoch) visualizer.print_logs('Average test PNSR is %f for %d images' % (avg_psnr/task_cnt, task_cnt))