def main(): args = parse_args() random.seed(0) torch.manual_seed(0) if not args.nogpu: torch.cuda.manual_seed_all(0) if args.no_norm: imgtr = [ToTensor()] else: imgtr = [ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [RandomSizedCrop((320, 320))] # (321,321) print("dataset_dir: ", args.dataset_dir) trainset_l = Corrosion(home_dir, args.dataset_dir, img_transform=Compose(imgtr), label_transform=Compose(labtr), co_transform=Compose(cotr), split=args.split, labeled=True) trainloader_l = DataLoader(trainset_l, batch_size=args.batch_size, shuffle=True, num_workers=2, drop_last=True) # if args.mode != 'semi': trainset_u = Corrosion(home_dir, args.dataset_dir, img_transform=Compose(imgtr), label_transform=Compose(labtr), co_transform=Compose(cotr), split=args.split, labeled=False) trainloader_u = DataLoader(trainset_u, batch_size=args.batch_size, shuffle=True, num_workers=2, drop_last=True) ######################### # Validation Dataloader # ######################## if args.val_orig: if args.no_norm: imgtr = [ZeroPadding(), ToTensor()] else: imgtr = [ZeroPadding(), ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [] else: if args.no_norm: imgtr = [ToTensor()] else: imgtr = [ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [RandomSizedCrop((320, 320))] # (321,321) valset = Corrosion(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \ label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False) valoader = DataLoader(valset, batch_size=1) ############# # GENERATOR # ############# # generator = deeplabv2.ResDeeplab() generator = unet.AttU_Net() # model_summary = generator.cuda() init_weights(generator, args.init_net) if args.init_net != 'unet': optimG = optim.SGD(filter(lambda p: p.requires_grad, \ generator.parameters()),lr=args.g_lr,momentum=0.9,\ weight_decay=0.0001,nesterov=True) else: optimG = optim.Adam(filter(lambda p: p.requires_grad, \ generator.parameters()),args.g_lr, [0.5, 0.999]) if not args.nogpu: generator = nn.DataParallel(generator).cuda() ################# # DISCRIMINATOR # ################ if args.mode != "base": discriminator = Dis(in_channels=2) # model_summary = discriminator.cuda() # summary(model_summary, (2, 320, 320)) if args.d_optim == 'adam': optimD = optim.Adam(filter(lambda p: p.requires_grad, \ discriminator.parameters()),lr = args.d_lr,weight_decay=0.0001) else: optimD = optim.SGD(filter(lambda p: p.requires_grad, \ discriminator.parameters()),lr=args.d_lr,weight_decay=0.0001,momentum=0.5,nesterov=True) if not args.nogpu: discriminator = nn.DataParallel(discriminator).cuda() if args.mode == 'base': train_base(generator, optimG, trainloader_l, valoader, args) elif args.mode == 'adv': train_adv(generator, discriminator, optimG, optimD, trainloader_l, valoader, args) elif args.mode == 'semi': train_semi(generator, discriminator, optimG, optimD, trainloader_l, trainloader_u, valoader, args) else: train_semir(generator, discriminator, optimG, optimD, trainloader_l, valoader, args)
def main(): home_dir = os.path.dirname(os.path.realpath(__file__)) dataset_dir = '/media/data/seg_dataset/corrosion/' test_img_list = os parser = argparse.ArgumentParser() parser.add_argument("dataset_dir", help="A directory containing img (Images) \ and cls (GT Segmentation) folder") parser.add_argument("snapshot", help="Snapshot with the saved model") parser.add_argument("--val_orig", help="Do Inference on original size image.\ Otherwise, crop to 321x321 like in training ", action='store_true') parser.add_argument("--norm",help="Normalize the test images",\ action='store_true') args = parser.parse_args() # print(args.val_orig, args.norm) if args.val_orig: img_transform = transforms.Compose([ToTensor()]) if args.norm: img_transform = transforms.Compose( [ToTensor(), NormalizeOwn(dataset='corrosion')]) label_transform = transforms.Compose( [IgnoreLabelClass(), ToTensorLabel()]) co_transform = transforms.Compose([RandomSizedCrop((321, 321))]) testset = Corrosion(home_dir, args.dataset_dir,img_transform=img_transform, \ label_transform = label_transform,co_transform=co_transform,train_phase=False) testloader = DataLoader(testset, batch_size=1) else: img_transform = transforms.Compose([ZeroPadding(), ToTensor()]) if args.norm: img_transform = img_transform = transforms.Compose( [ZeroPadding(), ToTensor(), NormalizeOwn(dataset='corrosion')]) label_transform = transforms.Compose( [IgnoreLabelClass(), ToTensorLabel()]) testset = Corrosion(home_dir,args.dataset_dir,img_transform=img_transform, \ label_transform=label_transform,train_phase=False) testloader = DataLoader(testset, batch_size=1) generator = deeplabv2.ResDeeplab() assert (os.path.isfile(args.snapshot)) snapshot = torch.load(args.snapshot) saved_net = { k.partition('module.')[2]: v for i, (k, v) in enumerate(snapshot['state_dict'].items()) } print('Snapshot Loaded') generator.load_state_dict(saved_net) generator.eval() generator = nn.DataParallel(generator).cuda() print('Generator Loaded') n_classes = 2 gts, preds = [], [] print('Prediction Goint to Start') # TODO: Crop out the padding before prediction for img_id, (img, gt_mask, _) in enumerate(testloader): print("Generating Predictions for Image {}".format(img_id)) gt_mask = gt_mask.numpy()[0] img = Variable(img.cuda()) out_pred_map = generator(img) # Get hard prediction soft_pred = out_pred_map.data.cpu().numpy()[0] soft_pred = soft_pred[:, :gt_mask.shape[0], :gt_mask.shape[1]] hard_pred = np.argmax(soft_pred, axis=0).astype(np.uint8) for gt_, pred_ in zip(gt_mask, hard_pred): gts.append(gt_) preds.append(pred_) score, class_iou = scores(gts, preds, n_class=n_classes) print("Mean IoU: {}".format(score))
def train_adv(args): if args.no_norm: imgtr = [ToTensor()] else: imgtr = [ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [RandomSizedCrop((321, 321))] trainset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), label_transform=Compose(labtr), \ co_transform=Compose(cotr)) trainloader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2, drop_last=True) ######################### # Validation Dataloader # ######################## if args.val_orig: if args.no_norm: imgtr = [ZeroPadding(), ToTensor()] else: imgtr = [ZeroPadding(), ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [] else: if args.no_norm: imgtr = [ToTensor()] else: imgtr = [ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [RandomSizedCrop((321, 321))] valset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \ label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False) valoader = DataLoader(valset, batch_size=1) ############# # GENERATOR # ############# generator = deeplabv2.ResDeeplab() optimG = optim.SGD(filter(lambda p: p.requires_grad, \ generator.parameters()),lr=args.g_lr,momentum=0.9,\ weight_decay=0.0001,nesterov=True) if not args.nogpu: generator = nn.DataParallel(generator).cuda() ################# # DISCRIMINATOR # ################ discriminator = Dis(in_channels=21) if args.d_optim == 'adam': optimD = optim.Adam(filter(lambda p: p.requires_grad, \ discriminator.parameters()),lr = args.d_lr) else: optimD = optim.SGD(filter(lambda p: p.requires_grad, \ discriminator.parameters()),lr=args.d_lr,weight_decay=0.0001,momentum=0.5,nesterov=True) if not args.nogpu: discriminator = nn.DataParallel(discriminator).cuda() ############# # TRAINING # ############# best_miou = -1 for epoch in range(args.start_epoch, args.max_epoch + 1): generator.train() for batch_id, (img, mask, ohmask) in enumerate(trainloader): if args.nogpu: img, mask, ohmask = Variable(img), Variable(mask), Variable( ohmask) else: img, mask, ohmask = Variable(img.cuda()), Variable( mask.cuda()), Variable(ohmask.cuda()) itr = len(trainloader) * (epoch - 1) + batch_id cpmap = generator(Variable(img.data, volatile=True)) cpmap = nn.Softmax2d()(cpmap) N = cpmap.size()[0] H = cpmap.size()[2] W = cpmap.size()[3] # Generate the Real and Fake Labels targetf = Variable(torch.zeros((N, H, W)).long(), requires_grad=False) targetr = Variable(torch.ones((N, H, W)).long(), requires_grad=False) if not args.nogpu: targetf = targetf.cuda() targetr = targetr.cuda() ########################## # DISCRIMINATOR TRAINING # ########################## optimD.zero_grad() # Train on Real confr = nn.LogSoftmax()(discriminator(ohmask.float())) if args.d_label_smooth != 0: LDr = (1 - args.d_label_smooth) * nn.NLLLoss2d()(confr, targetr) LDr += args.d_label_smooth * nn.NLLLoss2d()(confr, targetf) else: LDr = nn.NLLLoss2d()(confr, targetr) LDr.backward() # Train on Fake conff = nn.LogSoftmax()(discriminator(Variable(cpmap.data))) LDf = nn.NLLLoss2d()(conff, targetf) LDf.backward() poly_lr_scheduler(optimD, args.d_lr, itr) optimD.step() ###################### # GENERATOR TRAINING # ##################### optimG.zero_grad() cmap = generator(img) cpmapsmax = nn.Softmax2d()(cmap) cpmaplsmax = nn.LogSoftmax()(cmap) conff = nn.LogSoftmax()(discriminator(cpmapsmax)) LGce = nn.NLLLoss2d()(cpmaplsmax, mask) LGadv = nn.NLLLoss2d()(conff, targetr) LGseg = LGce + args.lam_adv * LGadv LGseg.backward() poly_lr_scheduler(optimG, args.g_lr, itr) optimG.step() print("[{}][{}] LD: {:.4f} LDfake: {:.4f} LD_real: {:.4f} LG: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f}" \ .format(epoch,itr,(LDr + LDf).data[0],LDr.data[0],LDf.data[0],LGseg.data[0],LGce.data[0],LGadv.data[0])) snapshot(generator, valoader, epoch, best_miou, args.snapshot_dir, args.prefix)
def train_semi(args): # TODO: Make it more generic to include for other splits args.batch_size = args.batch_size * 2 if args.no_norm: imgtr = [ToTensor()] else: imgtr = [ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [RandomSizedCrop((321, 321))] trainset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), label_transform=Compose(labtr), \ co_transform=Compose(cotr)) trainloader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2, drop_last=True) ######################### # Validation Dataloader # ######################## if args.val_orig: if args.no_norm: imgtr = [ZeroPadding(), ToTensor()] else: imgtr = [ZeroPadding(), ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [] else: if args.no_norm: imgtr = [ToTensor()] else: imgtr = [ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [RandomSizedCrop((321, 321))] valset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \ label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False) valoader = DataLoader(valset, batch_size=1) ############# # GENERATOR # ############# generator = deeplabv2.ResDeeplab() optimG = optim.SGD(filter(lambda p: p.requires_grad, \ generator.parameters()),lr=args.g_lr,momentum=0.9,\ weight_decay=0.0001,nesterov=True) if not args.nogpu: generator = nn.DataParallel(generator).cuda() ################# # DISCRIMINATOR # ################ discriminator = Dis(in_channels=21) if args.d_optim == 'adam': optimD = optim.Adam(filter(lambda p: p.requires_grad, \ discriminator.parameters()),lr = args.d_lr) else: optimD = optim.SGD(filter(lambda p: p.requires_grad, \ discriminator.parameters()),lr=args.d_lr,weight_decay=0.0001,momentum=0.5,nesterov=True) if not args.nogpu: discriminator = nn.DataParallel(discriminator).cuda() ############ # TRAINING # ############ best_miou = -1 for epoch in range(args.start_epoch, args.max_epoch + 1): generator.train() for batch_id, (img, mask, ohmask) in enumerate(trainloader): if args.nogpu: img, mask, ohmask = Variable(img), Variable(mask), Variable( ohmask) else: img, mask, ohmask = Variable(img.cuda()), Variable( mask.cuda()), Variable(ohmask.cuda()) itr = len(trainloader) * (epoch - 1) + batch_id ## TODO: Extend random interleaving for split of any size mid = args.batch_size // 2 img1, mask1, ohmask1 = img[0:mid, ...], mask[0:mid, ...], ohmask[0:mid, ...] img2, mask2, ohmask2 = img[mid:, ...], mask[mid:, ...], ohmask[mid:, ...] # Random Interleaving if random.random() < 0.5: imgl, maskl, ohmaskl = img1, mask1, ohmask1 imgu, masku, ohmasku = img2, mask2, ohmask2 else: imgu, masku, ohmasku = img1, mask1, ohmask1 imgl, maskl, ohmaskl = img2, mask2, ohmask2 ################################################ # Labelled data for Discriminator Training # ################################################ cpmap = generator(Variable(imgl.data, volatile=True)) cpmap = nn.Softmax2d()(cpmap) N = cpmap.size()[0] H = cpmap.size()[2] W = cpmap.size()[3] # Generate the Real and Fake Labels targetf = Variable(torch.zeros((N, H, W)).long()) targetr = Variable(torch.ones((N, H, W)).long()) if not args.nogpu: targetf = targetf.cuda() targetr = targetr.cuda() # Train on Real confr = nn.LogSoftmax()(discriminator(ohmaskl.float())) optimD.zero_grad() if args.d_label_smooth != 0: LDr = (1 - args.d_label_smooth) * nn.NLLLoss2d()(confr, targetr) LDr += args.d_label_smooth * nn.NLLLoss2d()(confr, targetf) else: LDr = nn.NLLLoss2d()(confr, targetr) LDr.backward() # Train on Fake conff = nn.LogSoftmax()(discriminator(Variable(cpmap.data))) LDf = nn.NLLLoss2d()(conff, targetf) LDf.backward() poly_lr_scheduler(optimD, args.d_lr, itr) optimD.step() ########################################### # labelled data Generator Training # ########################################### optimG.zero_grad() cpmap = generator(imgl) cpmapsmax = nn.Softmax2d()(cpmap) cpmaplsmax = nn.LogSoftmax()(cpmap) conff = nn.LogSoftmax()(discriminator(cpmapsmax)) LGce = nn.NLLLoss2d()(cpmaplsmax, maskl) LGadv = nn.NLLLoss2d()(conff, targetr) LGadv_d = LGadv.data[0] LGce_d = LGce.data[0] LGadv = args.lam_adv * LGadv (LGce + LGadv).backward() ##################################### # Use unlabelled data to get L_semi # ##################################### LGsemi_d = 0 if epoch > args.wait_semi: cpmap = generator(imgu) softpred = nn.Softmax2d()(cpmap) hardpred = torch.max(softpred, 1)[1].squeeze(1) conf = nn.Softmax2d()(discriminator( Variable(softpred.data, volatile=True))) idx = np.zeros(cpmap.data.cpu().numpy().shape, dtype=np.uint8) idx = idx.transpose(0, 2, 3, 1) confnp = cpmap[:, 1, ...].data.cpu().numpy() hardprednp = hardpred.data.cpu().numpy() idx[confnp > args.t_semi] = np.identity( 21, dtype=idx.dtype)[hardprednp[confnp > args.t_semi]] if np.count_nonzero(idx) != 0: cpmaplsmax = nn.LogSoftmax()(cpmap) idx = Variable(torch.from_numpy(idx).byte().cuda()) LGsemi_arr = cpmaplsmax.masked_select(idx) LGsemi = -1 * LGsemi_arr.mean() LGsemi_d = LGsemi.data[0] LGsemi = args.lam_semi * LGsemi LGsemi.backward() else: LGsemi_d = 0 LGseg_d = LGce_d + LGadv_d + LGsemi_d del idx del conf del confnp del hardpred del softpred del hardprednp del cpmap LGseg_d = LGce_d + LGadv_d + LGsemi_d poly_lr_scheduler(optimG, args.g_lr, itr) optimG.step() # Manually free memory! Later, really understand how computation graphs free variables print("[{}][{}] LD: {:.4f} LD_fake: {:.4f} LD_real: {:.4f} LG: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f} LG_semi: {:.4f}"\ .format(epoch,itr,(LDr + LDf).data[0],LDr.data[0],LDf.data[0],LGseg_d,LGce_d,LGadv_d,LGsemi_d)) snapshot(generator, valoader, epoch, best_miou, args.snapshot_dir, args.prefix)
def train_base(args): ####################### # Training Dataloader # ####################### if args.no_norm: imgtr = [ToTensor()] else: imgtr = [ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [RandomSizedCrop((321, 321))] trainset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), label_transform=Compose(labtr), \ co_transform=Compose(cotr)) trainloader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2, drop_last=True) ######################### # Validation Dataloader # ######################## if args.val_orig: if args.no_norm: imgtr = [ZeroPadding(), ToTensor()] else: imgtr = [ZeroPadding(), ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [] else: if args.no_norm: imgtr = [ToTensor()] else: imgtr = [ToTensor(), NormalizeOwn()] labtr = [IgnoreLabelClass(), ToTensorLabel()] cotr = [RandomSizedCrop((321, 321))] valset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \ label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False) valoader = DataLoader(valset, batch_size=1) model = deeplabv2.ResDeeplab() init_weights(model, args.init_net) optimG = optim.SGD(filter(lambda p: p.requires_grad, \ model.parameters()),lr=args.g_lr,momentum=0.9,\ weight_decay=0.0001,nesterov=True) if not args.nogpu: model = nn.DataParallel(model).cuda() best_miou = -1 for epoch in range(args.start_epoch, args.max_epoch + 1): model.train() for batch_id, (img, mask, _) in enumerate(trainloader): if args.nogpu: img, mask = Variable(img), Variable(mask) else: img, mask = Variable(img.cuda()), Variable(mask.cuda()) itr = len(trainloader) * (epoch - 1) + batch_id cprob = model(img) cprob = nn.LogSoftmax()(cprob) Lseg = nn.NLLLoss2d()(cprob, mask) poly_lr_scheduler(optimG, args.g_lr, itr) optimG.zero_grad() Lseg.backward() optimG.step() print("[{}][{}]Loss: {:0.4f}".format(epoch, itr, Lseg.data[0])) snapshot(model, valoader, epoch, best_miou, args.snapshot_dir, args.prefix)