def main(data_dir): # 0) Tensoboard Writer. writer = SummaryWriter(FLAGS['summary_path']) origin_img, uv_map_gt, uv_map_predicted = None, None, None if not os.path.exists(FLAGS['images']): os.mkdir(FLAGS['images']) # 1) Create Dataset of 300_WLP & Dataloader. wlp300 = PRNetDataset(root_dir=data_dir, transform=transforms.Compose([ ToTensor(), ToResize((416, 416)), ToNormalize(FLAGS["normalize_mean"], FLAGS["normalize_std"]) ])) wlp300_dataloader = DataLoader(dataset=wlp300, batch_size=FLAGS['batch_size'], shuffle=True, num_workers=1) # 2) Intermediate Processing. transform_img = transforms.Compose([ transforms.Normalize(FLAGS["normalize_mean"], FLAGS["normalize_std"]) ]) # 3) Create PRNet model. start_epoch, target_epoch = FLAGS['start_epoch'], FLAGS['target_epoch'] model = ResFCN256(resolution_input=416, resolution_output=416, channel=3, size=16) discriminator = Discriminator() # Load the pre-trained weight if FLAGS['resume'] != "" and os.path.exists( os.path.join(FLAGS['pretrained'], FLAGS['resume'])): state = torch.load(os.path.join(FLAGS['pretrained'], FLAGS['resume'])) model.load_state_dict(state['prnet']) start_epoch = state['start_epoch'] INFO("Load the pre-trained weight! Start from Epoch", start_epoch) else: start_epoch = 0 INFO( "Pre-trained weight cannot load successfully, train from scratch!") if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) model.to(FLAGS["device"]) discriminator.to(FLAGS["device"]) optimizer = torch.optim.Adam(model.parameters(), lr=FLAGS["lr"], betas=(0.5, 0.999)) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=FLAGS["lr"]) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99) stat_loss = SSIM(mask_path=FLAGS["mask_path"], gauss=FLAGS["gauss_kernel"]) loss = WeightMaskLoss(mask_path=FLAGS["mask_path"]) bce_loss = torch.nn.BCEWithLogitsLoss() bce_loss.to(FLAGS["device"]) #Loss function for adversarial for ep in range(start_epoch, target_epoch): bar = tqdm(wlp300_dataloader) loss_list_G, stat_list = [], [] loss_list_D = [] for i, sample in enumerate(bar): uv_map, origin = sample['uv_map'].to( FLAGS['device']), sample['origin'].to(FLAGS['device']) # Inference. optimizer.zero_grad() uv_map_result = model(origin) # Update D optimizer_D.zero_grad() fake_detach = uv_map_result.detach() d_fake = discriminator(fake_detach) d_real = discriminator(uv_map) retain_graph = False if FLAGS['gan_type'] == 'GAN': loss_d = bce_loss(d_real, d_fake) elif FLAGS['gan_type'].find('WGAN') >= 0: loss_d = (d_fake - d_real).mean() if FLAGS['gan_type'].find('GP') >= 0: epsilon = torch.rand(fake_detach.shape[0]).view( -1, 1, 1, 1) epsilon = epsilon.to(fake_detach.device) hat = fake_detach.mul(1 - epsilon) + uv_map.mul(epsilon) hat.requires_grad = True d_hat = discriminator(hat) gradients = torch.autograd.grad(outputs=d_hat.sum(), inputs=hat, retain_graph=True, create_graph=True, only_inputs=True)[0] gradients = gradients.view(gradients.size(0), -1) gradient_norm = gradients.norm(2, dim=1) gradient_penalty = 10 * gradient_norm.sub(1).pow(2).mean() loss_d += gradient_penalty # from ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks elif FLAGS['gan_type'] == 'RGAN': better_real = d_real - d_fake.mean(dim=0, keepdim=True) better_fake = d_fake - d_real.mean(dim=0, keepdim=True) loss_d = bce_loss(better_real, better_fake) retain_graph = True if discriminator.training: loss_list_D.append(loss_d.item()) loss_d.backward(retain_graph=retain_graph) optimizer_D.step() if 'WGAN' in FLAGS['gan_type']: for p in discriminator.parameters(): p.data.clamp_(-1, 1) # Update G d_fake_bp = discriminator( uv_map_result) # for backpropagation, use fake as it is if FLAGS['gan_type'] == 'GAN': label_real = torch.ones_like(d_fake_bp) loss_g = bce_loss(d_fake_bp, label_real) elif FLAGS['gan_type'].find('WGAN') >= 0: loss_g = -d_fake_bp.mean() elif FLAGS['gan_type'] == 'RGAN': better_real = d_real - d_fake_bp.mean(dim=0, keepdim=True) better_fake = d_fake_bp - d_real.mean(dim=0, keepdim=True) loss_g = bce_loss(better_fake, better_real) loss_g.backward() loss_list_G.append(loss_g.item()) optimizer.step() stat_logit = stat_loss(uv_map_result, uv_map) stat_list.append(stat_logit.item()) #bar.set_description(" {} [Loss(Paper)] {} [Loss(D)] {} [SSIM({})] {}".format(ep, loss_list_G[-1], loss_list_D[-1],FLAGS["gauss_kernel"], stat_list[-1])) # Record Training information in Tensorboard. """ if origin_img is None and uv_map_gt is None: origin_img, uv_map_gt = origin, uv_map uv_map_predicted = uv_map_result writer.add_scalar("Original Loss", loss_list_G[-1], FLAGS["summary_step"]) writer.add_scalar("D Loss", loss_list_D[-1], FLAGS["summary_step"]) writer.add_scalar("SSIM Loss", stat_list[-1], FLAGS["summary_step"]) grid_1, grid_2, grid_3 = make_grid(origin_img, normalize=True), make_grid(uv_map_gt), make_grid(uv_map_predicted) writer.add_image('original', grid_1, FLAGS["summary_step"]) writer.add_image('gt_uv_map', grid_2, FLAGS["summary_step"]) writer.add_image('predicted_uv_map', grid_3, FLAGS["summary_step"]) writer.add_graph(model, uv_map) """ if ep % FLAGS["save_interval"] == 0: with torch.no_grad(): print(" {} [Loss(Paper)] {} [Loss(D)] {} [SSIM({})] {}".format( ep, loss_list_G[-1], loss_list_D[-1], FLAGS["gauss_kernel"], stat_list[-1])) origin = cv2.imread("./test_data/obama_origin.jpg") gt_uv_map = np.load("./test_data/test_obama.npy") origin, gt_uv_map = test_data_preprocess( origin), test_data_preprocess(gt_uv_map) origin, gt_uv_map = transform_img(origin), transform_img( gt_uv_map) origin_in = origin.unsqueeze_(0).cuda() pred_uv_map = model(origin_in).detach().cpu() save_image( [origin.cpu(), gt_uv_map.unsqueeze_(0).cpu(), pred_uv_map], os.path.join(FLAGS['images'], str(ep) + '.png'), nrow=1, normalize=True) # Save model print("Save model") state = { 'prnet': model.state_dict(), 'Loss': loss_list_G, 'start_epoch': ep, 'Loss_D': loss_list_D, } torch.save(state, os.path.join(FLAGS['images'], '{}.pth'.format(ep))) scheduler.step() writer.close()
def main(data_dir): origin_img, uv_map_gt, uv_map_predicted = None, None, None if not os.path.exists(FLAGS['images']): os.mkdir(FLAGS['images']) # 1) Create Dataset of 300_WLP & Dataloader. wlp300 = PRNetDataset(root_dir=data_dir, transform=transforms.Compose([ ToTensor(), ToResize((416, 416)), ToNormalize(FLAGS["normalize_mean"], FLAGS["normalize_std"]) ])) wlp300_dataloader = DataLoader(dataset=wlp300, batch_size=FLAGS['batch_size'], shuffle=True, num_workers=1) # 2) Intermediate Processing. transform_img = transforms.Compose([ #transforms.ToTensor(), transforms.Normalize(FLAGS["normalize_mean"], FLAGS["normalize_std"]) ]) # 3) Create PRNet model. start_epoch, target_epoch = FLAGS['start_epoch'], FLAGS['target_epoch'] g_x = ResFCN256(resolution_input=416, resolution_output=416, channel=3, size=16) g_y = ResFCN256(resolution_input=416, resolution_output=416, channel=3, size=16) d_x = Discriminator() d_y = Discriminator() # Load the pre-trained weight if FLAGS['resume'] != "" and os.path.exists( os.path.join(FLAGS['pretrained'], FLAGS['resume'])): state = torch.load(os.path.join(FLAGS['pretrained'], FLAGS['resume'])) try: g_x.load_state_dict(state['g_x']) g_y.load_state_dict(state['g_y']) d_x.load_state_dict(state['d_x']) d_y.load_state_dict(state['d_y']) except Exception: g_x.load_state_dict(state['prnet']) start_epoch = state['start_epoch'] INFO("Load the pre-trained weight! Start from Epoch", start_epoch) else: start_epoch = 0 INFO( "Pre-trained weight cannot load successfully, train from scratch!") if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) g_x.to(FLAGS["device"]) g_y.to(FLAGS["device"]) d_x.to(FLAGS["device"]) d_y.to(FLAGS["device"]) optimizer_g = torch.optim.Adam(itertools.chain(g_x.parameters(), g_y.parameters()), lr=FLAGS["lr"], betas=(0.5, 0.999)) optimizer_d = torch.optim.Adam(itertools.chain(d_x.parameters(), d_y.parameters()), lr=FLAGS["lr"]) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer_g, gamma=0.99) stat_loss = SSIM(mask_path=FLAGS["mask_path"], gauss=FLAGS["gauss_kernel"]) loss = WeightMaskLoss(mask_path=FLAGS["mask_path"]) bce_loss = torch.nn.BCEWithLogitsLoss() bce_loss.to(FLAGS["device"]) l1_loss = nn.L1Loss().to(FLAGS["device"]) lambda_X = 10 lambda_Y = 10 #Loss function for adversarial for ep in range(start_epoch, target_epoch): bar = tqdm(wlp300_dataloader) loss_list_cycle_x = [] loss_list_cycle_y = [] loss_list_d_x = [] loss_list_d_y = [] real_label = torch.ones(FLAGS['batch_size']) fake_label = torch.zeros(FLAGS['batch_size']) for i, sample in enumerate(bar): real_y, real_x = sample['uv_map'].to( FLAGS['device']), sample['origin'].to(FLAGS['device']) # x -> y' -> x^ optimizer_g.zero_grad() fake_y = g_x(real_x) prediction = d_x(fake_y) loss_g_x = bce_loss(prediction, real_label) x_hat = g_y(fake_y) loss_cycle_x = l1_loss(x_hat, real_x) * lambda_X loss_x = loss_g_x + loss_cycle_x loss_x.backward(retain_graph=True) optimizer_g.step() loss_list_cycle_x.append(loss_x.item()) # y -> x' -> y^ optimizer_g.zero_grad() fake_x = g_y(real_y) prediction = d_y(fake_x) loss_g_y = bce_loss(prediction, real_label) y_hat = g_x(fake_x) loss_cycle_y = l1_loss(y_hat, real_y) * lambda_Y loss_y = loss_g_y + loss_cycle_y loss_y.backward(retain_graph=True) optimizer_g.step() loss_list_cycle_y.append(loss_y.item()) # d_x optimizer_d.zero_grad() pred_real = d_x(real_y) loss_d_x_real = bce_loss(pred_real, real_label) pred_fake = d_x(fake_y) loss_d_x_fake = bce_loss(pred_fake, fake_label) loss_d_x = (loss_d_x_real + loss_d_x_fake) * 0.5 loss_d_x.backward() loss_list_d_x.append(loss_d_x.item()) optimizer_d.step() if 'WGAN' in FLAGS['gan_type']: for p in d_x.parameters(): p.data.clamp_(-1, 1) # d_y optimizer_d.zero_grad() pred_real = d_y(real_x) loss_d_y_real = bce_loss(pred_real, real_label) pred_fake = d_y(fake_x) loss_d_y_fake = bce_loss(pred_fake, fake_label) loss_d_y = (loss_d_y_real + loss_d_y_fake) * 0.5 loss_d_y.backward() loss_list_d_y.append(loss_d_y.item()) optimizer_d.step() if 'WGAN' in FLAGS['gan_type']: for p in d_y.parameters(): p.data.clamp_(-1, 1) if ep % FLAGS["save_interval"] == 0: with torch.no_grad(): print( " {} [Loss_G_X] {} [Loss_G_Y] {} [Loss_D_X] {} [Loss_D_Y] {}" .format(ep, loss_list_g_x[-1], loss_list_g_y[-1], loss_list_d_x[-1], loss_list_d_y[-1])) origin = cv2.imread("./test_data/obama_origin.jpg") gt_uv_map = np.load("./test_data/test_obama.npy") origin, gt_uv_map = test_data_preprocess( origin), test_data_preprocess(gt_uv_map) origin, gt_uv_map = transform_img(origin), transform_img( gt_uv_map) origin_in = origin.unsqueeze_(0).cuda() pred_uv_map = g_x(origin_in).detach().cpu() save_image( [origin.cpu(), gt_uv_map.unsqueeze_(0).cpu(), pred_uv_map], os.path.join(FLAGS['images'], str(ep) + '.png'), nrow=1, normalize=True) # Save model print("Save model") state = { 'g_x': g_x.state_dict(), 'g_y': g_y.state_dict(), 'd_x': d_x.state_dict(), 'd_y': d_y.state_dict(), 'start_epoch': ep, } torch.save(state, os.path.join(FLAGS['images'], '{}.pth'.format(ep))) scheduler.step()
for each_hazy in hazy_lst: hazy_list.append(hazy_img_path + each_hazy) clear_list.append(clear_img_path + each_hazy.split('_')[0] + '.png') #divide the dataset into train/val train_loader, test_loader = split_train_val(train_size, image_size, hazy_list, clear_list, batch_size) #inititialize the Generator,Discriminator and move them to GPU if available gen = Generator() gen.apply(init_weights) gen.to(device) disc = Discriminator() disc.apply(init_weights) disc.to(device) #initialize the gan,content,perceptual and brightness loss gan_criterion = nn.BCELoss().to(device) content_criterion = nn.L1Loss().to(device) perceptual_criterion = nn.MSELoss().to(device) brightness_criterion = nn.L1Loss().to(device) #set the feature extractor to evaluation mode as it will be used only to calculate perceptual loss feature_extractor = FeatureExtractor() feature_extractor.eval() feature_extractor.to(device) #initialize the optimizers for Generator and Discriminator optimizerD = optim.Adam(disc.parameters(), lr=0.0003, betas=(0.5, 0.999)) optimizerG = optim.Adam(gen.parameters(), lr=0.0001, betas=(0.5, 0.999))
def train(train_sources, eval_source): path = sys.argv[1] dr = DataReader(path, train_sources) dr.read() print(len(dr.train.x)) batch_size = 8 device = torch.device('cpu') if torch.cuda.is_available(): device = torch.device('cuda') dataset_s_train = MultiDomainDataset(dr.train.x, dr.train.y, dr.train.vendor, device, DomainAugmentation()) dataset_s_dev = MultiDomainDataset(dr.dev.x, dr.dev.y, dr.dev.vendor, device) dataset_s_test = MultiDomainDataset(dr.test.x, dr.test.y, dr.test.vendor, device) loader_s_train = DataLoader(dataset_s_train, batch_size, shuffle=True) dr_eval = DataReader(path, [eval_source]) dr_eval.read() dataset_eval_dev = MultiDomainDataset(dr_eval.dev.x, dr_eval.dev.y, dr_eval.dev.vendor, device) dataset_eval_test = MultiDomainDataset(dr_eval.test.x, dr_eval.test.y, dr_eval.test.vendor, device) dataset_da_train = MultiDomainDataset(dr.train.x+dr_eval.train.x, dr.train.y+dr_eval.train.y, dr.train.vendor+dr_eval.train.vendor, device, DomainAugmentation()) loader_da_train = DataLoader(dataset_da_train, batch_size, shuffle=True) segmentator = UNet() discriminator = Discriminator(n_domains=len(train_sources)) discriminator.to(device) segmentator.to(device) sigmoid = nn.Sigmoid() selector = Selector() s_criterion = nn.BCELoss() d_criterion = nn.CrossEntropyLoss() s_optimizer = optim.AdamW(segmentator.parameters(), lr=0.0001, weight_decay=0.01) d_optimizer = optim.AdamW(discriminator.parameters(), lr=0.001, weight_decay=0.01) a_optimizer = optim.AdamW(segmentator.encoder.parameters(), lr=0.001, weight_decay=0.01) lmbd = 1/150 s_train_losses = [] s_dev_losses = [] d_train_losses = [] eval_domain_losses = [] train_dices = [] dev_dices = [] eval_dices = [] epochs = 3 da_loader_iter = iter(loader_da_train) for epoch in tqdm(range(epochs)): s_train_loss = 0.0 d_train_loss = 0.0 for index, sample in enumerate(loader_s_train): img = sample['image'] target_mask = sample['target'] da_sample = next(da_loader_iter, None) if epoch == 100: s_optimizer.defaults['lr'] = 0.001 d_optimizer.defaults['lr'] = 0.0001 if da_sample is None: da_loader_iter = iter(loader_da_train) da_sample = next(da_loader_iter, None) if epoch < 50 or epoch >= 100: # Training step of segmentator predicted_activations, inner_repr = segmentator(img) predicted_mask = sigmoid(predicted_activations) s_loss = s_criterion(predicted_mask, target_mask) s_optimizer.zero_grad() s_loss.backward() s_optimizer.step() s_train_loss += s_loss.cpu().detach().numpy() if epoch >= 50: # Training step of discriminator predicted_activations, inner_repr = segmentator(da_sample['image']) predicted_activations = predicted_activations.clone().detach() inner_repr = inner_repr.clone().detach() predicted_vendor = discriminator(predicted_activations, inner_repr) d_loss = d_criterion(predicted_vendor, da_sample['vendor']) d_optimizer.zero_grad() d_loss.backward() d_optimizer.step() d_train_loss += d_loss.cpu().detach().numpy() if epoch >= 100: # adversarial training step predicted_mask, inner_repr = segmentator(da_sample['image']) predicted_vendor = discriminator(predicted_mask, inner_repr) a_loss = -1 * lmbd * d_criterion(predicted_vendor, da_sample['vendor']) a_optimizer.zero_grad() a_loss.backward() a_optimizer.step() lmbd += 1/150 inference_model = nn.Sequential(segmentator, selector, sigmoid) inference_model.to(device) inference_model.eval() d_train_losses.append(d_train_loss / len(loader_s_train)) s_train_losses.append(s_train_loss / len(loader_s_train)) s_dev_losses.append(calculate_loss(dataset_s_dev, inference_model, s_criterion, batch_size)) eval_domain_losses.append(calculate_loss(dataset_eval_dev, inference_model, s_criterion, batch_size)) train_dices.append(calculate_dice(inference_model, dataset_s_train)) dev_dices.append(calculate_dice(inference_model, dataset_s_dev)) eval_dices.append(calculate_dice(inference_model, dataset_eval_dev)) segmentator.train() date_time = datetime.now().strftime("%m%d%Y_%H%M%S") model_path = os.path.join(pathlib.Path(__file__).parent.absolute(), "model", "weights", "segmentator"+str(date_time)+".pth") torch.save(segmentator.state_dict(), model_path) util.plot_data([(s_train_losses, 'train_losses'), (s_dev_losses, 'dev_losses'), (d_train_losses, 'discriminator_losses'), (eval_domain_losses, 'eval_domain_losses')], 'losses.png') util.plot_dice([(train_dices, 'train_dice'), (dev_dices, 'dev_dice'), (eval_dices, 'eval_dice')], 'dices.png') inference_model = nn.Sequential(segmentator, selector, sigmoid) inference_model.to(device) inference_model.eval() print('Dice on annotated: ', calculate_dice(inference_model, dataset_s_test)) print('Dice on unannotated: ', calculate_dice(inference_model, dataset_eval_test))