def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device train_logger = make_logger("Train.log", args) test_logger = make_logger("Test.log", args) model = load_models( mode="cls", device=device, args=args, ) optimizer = optim.Adam( model.parameters(), lr=args.lr, betas=(0.9, 0.999), ) optimizer.zero_grad() if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None if (args.train or args.run_semi) and args.test: print("===================================") print("====== Loading Training Data ======") print("===================================") sample_gt_list = np.load(args.gt_sample_list) trainset_gt = ModelNetDatasetGT( root_list=args.train_file, sample_list=sample_gt_list, ) trainloader_gt = torch.utils.data.DataLoader( trainset_gt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) print("===================================") print("====== Loading Test Data ======") print("===================================") testset = ModelNetDatasetGT( root_list=args.test_file, sample_list=None, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) args.iter_per_epoch = int(trainset_gt.__len__() / args.batch_size) args.total_data = trainset_gt.__len__() args.total_iterations = int(args.num_epochs * args.total_data / args.batch_size) args.iter_save_epoch = args.save_per_epoch * int( args.total_data / args.batch_size) args.iter_test_epoch = args.test_epoch * int( args.total_data / args.batch_size) if args.train and args.test: model.train() model.to(args.device) cls_loss = torch.nn.CrossEntropyLoss().to(device) trainloader_gt_iter = enumerate(trainloader_gt) run_training_pointnet_cls( trainloader_gt=trainloader_gt, trainloader_gt_iter=trainloader_gt_iter, testloader=testloader, model=model, cls_loss=cls_loss, optimizer=optimizer, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) if args.test: print("===================================") print("====== Loading Testing Data =======") print("===================================") testset = ModelNetDatasetGT( root_list=args.test_file, sample_list=None, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) criterion = torch.nn.CrossEntropyLoss().to(device) run_testing( dataloader=testloader, model=model, criterion=criterion, logger=test_logger, test_iter=100000000, writer=None, args=args, )
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) print("EXP PATH: {}".format(args.exp_dir)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device logger = make_logger(filename="TrainVal.log", args=args) if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None print("===================================") print("====== Loading Training Data ======") print("===================================") shape_transformer = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) photo_transformer = PhotometricTransform(photometric_transform_config) source_data = UnityDataset( root=args.source_root, image_size=args.image_size, shape_transforms=shape_transformer, photo_transforms=photo_transformer, train_bool=False, ) args.tot_source = source_data.__len__() args.total_iterations = args.num_epochs * args.tot_source // args.batch_size args.iters_to_eval = args.epoch_to_eval * args.tot_source // args.batch_size train_loader = torch.utils.data.DataLoader( source_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) model_seg = load_models( mode="segmentation", device=device, args=args, ) optimizer_seg = optim.Adam( model_seg.parameters(), lr=args.lr_seg, betas=(args.beta1, 0.999), ) optimizer_seg.zero_grad() seg_loss_target = torch.nn.CrossEntropyLoss().to(device) trainloader_iter = enumerate(train_loader) loss_seg_min = float("inf") miou_max = float("-inf") for i_iter in range(args.total_iterations): loss_seg_value = 0 model_seg.train() optimizer_seg.zero_grad() adjust_learning_rate( optimizer=optimizer_seg, learning_rate=args.lr_seg, i_iter=i_iter, max_steps=args.total_iterations, power=0.9, ) try: _, batch = next(trainloader_iter) except StopIteration: trainloader_iter = enumerate(train_loader) _, batch = next(trainloader_iter) images, labels = batch images = Variable(images).to(args.device) labels = Variable(labels.long()).to(args.device) pred = model_seg(images) loss_seg = seg_loss_target(pred, labels) current_loss_seg = loss_seg.item() loss_seg_value += current_loss_seg loss_seg.backward() optimizer_seg.step() pred_img = pred.argmax(dim=1, keepdim=True) flat_pred = pred_img.detach().cpu().numpy().flatten() flat_gt = labels.detach().cpu().numpy().flatten() miou, _ = compute_mean_iou(flat_pred=flat_pred, flat_label=flat_gt) logger.info('iter = {0:8d}/{1:8d} ' 'loss_seg = {2:.3f} ' 'mIoU = {3:.3f} '.format( i_iter, args.total_iterations, loss_seg_value, miou, )) if args.tensorboard and (writer != None): writer.add_scalar('Train/Cross_Entropy', current_loss_seg, i_iter) writer.add_scalar('Train/mIoU', miou, i_iter) if i_iter % args.iters_to_eval == 0: filename = os.path.join( args.exp_dir, "Target_img_trainiter_{}.png".format(i_iter)) target_img = images.float() gen_target_img = torchvision.utils.make_grid(target_img, padding=2, normalize=True) torchvision.utils.save_image(gen_target_img, filename) filename = os.path.join( args.exp_dir, "Unity_pred_trainiter_{}.png".format(i_iter)) pred_img = pred_img.float() gen_img = torchvision.utils.make_grid(pred_img, padding=2, normalize=True) torchvision.utils.save_image(gen_img, filename) is_better_ss = current_loss_seg < loss_seg_min if is_better_ss: loss_seg_min = current_loss_seg torch.save(model_seg.state_dict(), os.path.join(args.exp_dir, "model_train_best.pth")) if miou > miou_max: miou_max = miou torch.save(model_seg.state_dict(), os.path.join(args.exp_dir, "model_train_best_miou.pth")) logger.info("==========================================") logger.info("Training DONE!") if args.tensorboard and (writer != None): writer.close()
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device train_logger = make_logger("Train.log", args) test_logger = make_logger("Test.log", args) model = load_models( mode="seg", device=device, args=args, ) optimizer = optim.Adam( model.parameters(), lr=args.lr, betas=(0.9, 0.999), ) optimizer.zero_grad() if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None if args.train and args.test: print("===================================") print("====== Loading Training Data ======") print("===================================") if args.gt_sample_list != None: sample_gt_list = np.load(args.gt_sample_list) else: sample_gt_list = None trainset_gt = ShapeNetDatasetGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainloader_gt = torch.utils.data.DataLoader( trainset_gt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) print("===================================") print("====== Loading Test Data ======") print("===================================") testset = ShapeNetDatasetGT( root_list=args.test_file, sample_list=None, num_classes=16, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) args.iter_per_epoch = int(trainset_gt.__len__() / args.batch_size) args.total_data = trainset_gt.__len__() args.total_iterations = int(args.num_epochs * args.total_data / args.batch_size) args.iter_save_epoch = args.save_per_epoch * int( args.total_data / args.batch_size) args.iter_test_epoch = args.test_epoch * int( args.total_data / args.batch_size) if args.train and args.test: model.train() model.to(args.device) seg_loss = torch.nn.CrossEntropyLoss().to(device) trainloader_gt_iter = enumerate(trainloader_gt) run_training_pointnet_seg( trainloader_gt=trainloader_gt, trainloader_gt_iter=trainloader_gt_iter, testloader=testloader, testdataset=testset, model=model, seg_loss=seg_loss, optimizer=optimizer, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) if args.test: print("===================================") print("====== Loading Testing Data =======") print("===================================") testset = ShapeNetDatasetGT( root_list=args.test_file, sample_list=None, num_classes=16, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) criterion = torch.nn.CrossEntropyLoss().to(device) run_testing_seg( dataloader=testloader, dataset=testset, model=model, criterion=criterion, logger=test_logger, test_iter=100000000, writer=None, args=args, ) if args.tsne: from utils.metric import batch_get_iou, object_names from torch.autograd import Variable args.batch_size = 1 labels = [] objects = [] if args.gt_sample_list != None: sample_gt_list = np.load(args.gt_sample_list) else: sample_gt_list = None trainset_gt = ShapeNetDatasetGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainset_nogt = ShapeNetDataset_noGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainloader_gt = torch.utils.data.DataLoader( trainset_gt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) trainloader_nogt = torch.utils.data.DataLoader( trainset_nogt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) # print("===================================") # print("====== Loading Test Data ======") # print("===================================") # testset = ShapeNetDatasetGT( # root_list=args.test_file, # sample_list=None, # num_classes=16, # ) # testloader = torch.utils.data.DataLoader( # testset, # batch_size=args.batch_size, # shuffle=False, # num_workers=args.workers, # pin_memory=True, # ) # # print("===================================") # print("====== Loading Testing Data =======") # print("===================================") # testset = ShapeNetDatasetGT( # root_list=args.test_file, # sample_list=None, # num_classes=16, # ) # testloader = torch.utils.data.DataLoader( # testset, # batch_size=args.batch_size, # shuffle=False, # num_workers=args.workers, # pin_memory=True, # ) model.eval() shape_ious = np.empty(len(object_names), dtype=np.object) for i in range(shape_ious.shape[0]): shape_ious[i] = [] all_shapes_train_gt = np.empty((len(trainset_gt), 2048)) for batch_idx, data in enumerate(trainloader_gt): if batch_idx % 1000 == 0: print("Processing {} ...".format(batch_idx)) pts, cls, seg = data pts, cls, seg = Variable(pts).float(), \ Variable(cls), Variable(seg).type(torch.LongTensor) pts, cls, seg = pts.to(args.device), cls.to( args.device), seg.long().to(args.device) labels.append(1) objects.append(int(cls.argmax(axis=2).squeeze().cpu().numpy())) with torch.set_grad_enabled(False): pred, global_shape = model(pts, cls) all_shapes_train_gt[ batch_idx, :] = global_shape.squeeze().detach().cpu().numpy() all_shapes_train_nogt = np.empty((len(trainset_nogt), 2048)) for batch_idx, data in enumerate(trainloader_nogt): if batch_idx % 1000 == 0: print("Processing {} ...".format(batch_idx)) pts, cls = data pts, cls = Variable(pts).float(), Variable(cls) pts, cls = pts.to(args.device), cls.to(args.device) with torch.set_grad_enabled(False): pred, global_shape = model(pts, cls) all_shapes_train_nogt[ batch_idx, :] = global_shape.squeeze().detach().cpu().numpy() labels.append(0) objects.append(int(cls.argmax(axis=2).squeeze().cpu().numpy())) all_shapes = np.concatenate( (all_shapes_train_gt, all_shapes_train_nogt), axis=0) shape_info = { "shapes": all_shapes, "labels": labels, "objects": objects } import pickle try: o = open("global_shape_{}.pkl".format(len(trainset_gt)), "wb") pickle.dump(shape_info, o, protocol=2) o.close() except FileNotFoundError as e: print(e)
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) print("EXP PATH: {}".format(args.exp_dir)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device logger = make_logger(filename="TrainVal.log", args=args) if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None print("===================================") print("====== Loading Training Data ======") print("===================================") shape_transformer = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) photo_transformer = PhotometricTransform(photometric_transform_config) joint_data = JointDataset( root_source=args.source_root, image_size=args.image_size, data_to_train="dataloaders/eye/trinity_train_200.pkl", shape_transforms=shape_transformer, photo_transforms=photo_transformer, train_bool=False, ) args.tot_data = joint_data.__len__() args.total_iterations = args.num_epochs * args.tot_data // args.batch_size args.iters_to_eval = args.epoch_to_eval * args.tot_data // args.batch_size print("===================================") print("========= Loading Val Data ========") print("===================================") val_target_data = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "validation"), image_size=args.image_size, data_to_train="", shape_transforms=shape_transformer, photo_transforms=None, train_bool=False, ) train_loader = torch.utils.data.DataLoader( joint_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) val_loader = torch.utils.data.DataLoader( val_target_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) model_seg = load_models( mode="segmentation", device=device, args=args, ) optimizer_seg = optim.Adam( model_seg.parameters(), lr=args.lr_seg, betas=(args.beta1, 0.999), ) optimizer_seg.zero_grad() # class_weight_target = 1.0 / train_target_data.get_class_probability().to(device) seg_loss_target = torch.nn.CrossEntropyLoss().to(device) trainloader_iter = enumerate(train_loader) val_loss, val_miou = [], [] val_loss_f = float("inf") val_miou_f = float("-inf") loss_seg_min = float("inf") for i_iter in range(args.total_iterations): loss_seg_value = 0 model_seg.train() optimizer_seg.zero_grad() adjust_learning_rate( optimizer=optimizer_seg, learning_rate=args.lr_seg, i_iter=i_iter, max_steps=args.total_iterations, power=0.9, ) try: _, batch = next(trainloader_iter) except StopIteration: trainloader_iter = enumerate(train_loader) _, batch = next(trainloader_iter) images, labels = batch images = Variable(images).to(args.device) labels = Variable(labels.long()).to(args.device) pred = model_seg(images) loss_seg = seg_loss_target(pred, labels) current_loss_seg = loss_seg.item() loss_seg_value += current_loss_seg loss_seg.backward() optimizer_seg.step() logger.info('iter = {0:8d}/{1:8d} ' 'loss_seg = {2:.3f} '.format( i_iter, args.total_iterations, loss_seg_value, )) current_epoch = i_iter * args.batch_size // args.tot_data if i_iter % args.iters_to_eval == 0: val_loss_f, val_miou_f = validate_baseline( i_iter=i_iter, val_loader=val_loader, model=model_seg, epoch=current_epoch, logger=logger, writer=writer, val_loss=val_loss_f, val_iou=val_miou_f, args=args, ) val_loss.append(val_loss_f) val_loss_f = np.min(np.array(val_loss)) val_miou.append(val_miou_f) val_miou_f = np.max(np.array(val_miou)) if args.tensorboard and (writer != None): writer.add_scalar('Val/Cross_Entropy_Target', val_loss_f, i_iter) writer.add_scalar('Val/mIoU_Target', val_miou_f, i_iter) is_better_ss = current_loss_seg < loss_seg_min if is_better_ss: loss_seg_min = current_loss_seg torch.save(model_seg.state_dict(), os.path.join(args.exp_dir, "model_train_best.pth")) logger.info("==========================================") logger.info("Training DONE!") if args.tensorboard and (writer != None): writer.close() with open("%s/train_performance.pkl" % args.exp_dir, "wb") as f: pickle.dump([val_loss, val_miou], f) logger.info("==========================================") logger.info("Evaluating on test data ...") testdata = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "test"), image_size=args.image_size, data_to_train="", shape_transforms=shape_transformer, photo_transforms=None, train_bool=False, ) test_loader = torch.utils.data.DataLoader( testdata, batch_size=1, shuffle=False, num_workers=0, pin_memory=True, ) pm = run_testing( dataset=testdata, test_loader=test_loader, model=model_seg, args=args, ) logger.info('Global Mean Accuracy: {:.3f}'.format(np.array(pm.GA).mean())) logger.info('Mean IOU: {:.3f}'.format(np.array(pm.IOU).mean())) logger.info('Mean Recall: {:.3f}'.format(np.array(pm.Recall).mean())) logger.info('Mean Precision: {:.3f}'.format(np.array(pm.Precision).mean())) logger.info('Mean F1: {:.3f}'.format(np.array(pm.F1).mean())) IOU_ALL = np.array(pm.Iou_all) logger.info( "Back: {:.4f}, Sclera: {:.4f}, Iris: {:.4f}, Pupil: {:.4f}".format( IOU_ALL[:, 0].mean(), IOU_ALL[:, 1].mean(), IOU_ALL[:, 2].mean(), IOU_ALL[:, 3].mean(), ))
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device train_logger = make_logger("Train.log", args) test_logger = make_logger("Test.log", args) model = load_models( mode="cls", device=device, args=args, ) model_D = load_models( mode="disc", device=device, args=args, ) optimizer = optim.Adam( model.parameters(), lr=args.lr, betas=(0.9, 0.999), ) optimizer.zero_grad() optimizer_D = optim.Adam( model_D.parameters(), lr=args.lr_D, betas=(0.9, 0.999), ) optimizer_D.zero_grad() if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None if (args.train or args.run_semi) and args.test: print("===================================") print("====== Loading Training Data ======") print("===================================") # idx = np.arange(9840) # np.random.shuffle(idx) # sample_gt_list = idx[0:args.num_samples] # sample_nogt_list = idx[args.num_samples:] # filename = "gt_sample_{}.npy".format(args.name) # np.save(filename, sample_gt_list) sample_gt_list = np.load(args.gt_sample_list) trainset_gt = ModelNetDatasetGT( root_list=args.train_file, sample_list=sample_gt_list, ) trainset_nogt = ModelNetDataset_noGT( root_list=args.train_file, sample_list=sample_gt_list, ) trainloader_gt = torch.utils.data.DataLoader( trainset_gt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) trainloader_nogt = torch.utils.data.DataLoader( trainset_nogt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) print("===================================") print("====== Loading Test Data ======") print("===================================") testset = ModelNetDatasetGT( root_list=args.test_file, sample_list=None, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) args.iter_per_epoch = int(1.0 * trainset_gt.__len__() / args.batch_size) args.total_data = trainset_gt.__len__() + trainset_nogt.__len__() args.total_iterations = int(args.num_epochs * args.total_data / args.batch_size) args.iter_save_epoch = args.save_per_epoch * int(args.total_data / args.batch_size) args.iter_test_epoch = args.test_epoch * int(args.total_data / args.batch_size) args.semi_start = int(args.semi_start_epoch * args.total_data / args.batch_size) if args.train and args.test: model.train() model_D.train() model.to(args.device) model_D.to(args.device) #class_weight = 1.0 / trainset_gt.get_class_probability().to(device) cls_loss = torch.nn.CrossEntropyLoss().to(device) gan_loss = torch.nn.BCEWithLogitsLoss().to(device) semi_loss = torch.nn.CrossEntropyLoss(ignore_index=255) history_pool_gt = ImagePool(args.pool_size) history_pool_nogt = ImagePool(args.pool_size) trainloader_gt_iter = enumerate(trainloader_gt) targetloader_nogt_iter = enumerate(trainloader_nogt) if args.semi_start_epoch==0: run_training( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, model=model, model_D=model_D, gan_loss=gan_loss, cls_loss=cls_loss, optimizer=optimizer, optimizer_D=optimizer_D, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) else: run_training_semi( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, model=model, model_D=model_D, gan_loss=gan_loss, cls_loss=cls_loss, semi_loss=semi_loss, optimizer=optimizer, optimizer_D=optimizer_D, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) if args.test: print("===================================") print("====== Loading Testing Data =======") print("===================================") testset = ModelNetDatasetGT( root_list=args.test_file, sample_list=None, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) criterion = torch.nn.CrossEntropyLoss().to(device) run_testing( dataloader=testloader, model=model, criterion=criterion, logger=test_logger, test_iter=100000000, writer=None, args=args, )
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) print("EXP PATH: {}".format(args.exp_dir)) assert args.trinity_data_train_with_labels != "", "Indicate trained data in Trinity!" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device logger = make_logger(filename="TrainVal.log", args=args) if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None print("===================================") print("====== Loading Training Data ======") print("===================================") transforms_shape_source = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) transforms_shape_target = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) photo_transformer = PhotometricTransform(photometric_transform_config) source_data = UnityDataset( root=args.source_root, image_size=args.image_size, shape_transforms=transforms_shape_source, photo_transforms=photo_transformer, train_bool=False, ) target_data = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "train"), image_size=args.image_size, data_to_train=args.trinity_data_train_with_labels, shape_transforms=transforms_shape_target, photo_transforms=photo_transformer, train_bool=False, ) args.tot_source = source_data.__len__() args.total_iterations = args.num_epochs * source_data.__len__( ) // args.batch_size args.iters_to_eval = args.epoch_to_eval * source_data.__len__( ) // args.batch_size args.iter_source_to_eval = args.epoch_to_eval_source * source_data.__len__( ) // args.batch_size args.iter_semi_start = args.epoch_semi_start * source_data.__len__( ) // args.batch_size print("===================================") print("========= Loading Val Data ========") print("===================================") val_target_data = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "validation"), image_size=args.image_size, data_to_train="", shape_transforms=transforms_shape_target, photo_transforms=None, train_bool=False, ) # class_weight_source = 1.0 / source_data.get_class_probability().to(device) source_loader = torch.utils.data.DataLoader( source_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) target_loader = torch.utils.data.DataLoader( target_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) val_loader = torch.utils.data.DataLoader( val_target_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) model_seg = load_models( mode="segmentation", device=device, args=args, ) model_disc = load_models( mode="single_discriminator", device=device, args=args, ) optimizer_seg = optim.Adam( model_seg.parameters(), lr=args.lr_seg, betas=(args.beta1, 0.999), ) optimizer_seg.zero_grad() optimizer_disc = torch.optim.Adam( model_disc.parameters(), lr=args.lr_disc, betas=(args.beta1, 0.999), ) optimizer_disc.zero_grad() seg_loss_source = torch.nn.CrossEntropyLoss().to(device) gan_loss = torch.nn.BCEWithLogitsLoss().to(device) semi_loss = torch.nn.CrossEntropyLoss(ignore_index=-1) history_true_mask = ImagePool(args.pool_size) history_fake_mask = ImagePool(args.pool_size) trainloader_iter = enumerate(source_loader) targetloader_iter = enumerate(target_loader) val_loss, val_miou = run_training_SSDA( trainloader_source=source_loader, trainloader_target=target_loader, trainloader_iter=trainloader_iter, targetloader_iter=targetloader_iter, val_loader=val_loader, model_seg=model_seg, model_disc=model_disc, gan_loss=gan_loss, seg_loss=seg_loss_source, semi_loss_criterion=semi_loss, optimizer_seg=optimizer_seg, optimizer_disc=optimizer_disc, history_pool_true=history_true_mask, history_pool_fake=history_fake_mask, logger=logger, writer=writer, args=args, ) with open("%s/train_performance.pkl" % args.exp_dir, "wb") as f: pickle.dump([val_loss, val_miou], f) logger.info("==========================================") logger.info("Evaluating on test data ...") testdata = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "test"), image_size=args.image_size, data_to_train="", shape_transforms=transforms_shape_target, photo_transforms=None, train_bool=False, ) test_loader = torch.utils.data.DataLoader( testdata, batch_size=1, shuffle=False, num_workers=0, pin_memory=True, ) pm = run_testing( dataset=testdata, test_loader=test_loader, model=model_seg, args=args, ) logger.info('Global Mean Accuracy: {:.3f}'.format(np.array(pm.GA).mean())) logger.info('Mean IOU: {:.3f}'.format(np.array(pm.IOU).mean())) logger.info('Mean Recall: {:.3f}'.format(np.array(pm.Recall).mean())) logger.info('Mean Precision: {:.3f}'.format(np.array(pm.Precision).mean())) logger.info('Mean F1: {:.3f}'.format(np.array(pm.F1).mean())) IOU_ALL = np.array(pm.Iou_all) logger.info( "Back: {:.4f}, Sclera: {:.4f}, Iris: {:.4f}, Pupil: {:.4f}".format( IOU_ALL[:, 0].mean(), IOU_ALL[:, 1].mean(), IOU_ALL[:, 2].mean(), IOU_ALL[:, 3].mean(), ))
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device test_logger = make_logger("Test.log", args) model = load_models( mode="cls", device=device, args=args, ) # model_D = load_models( # mode="disc", # device=device, # args=args, # ) optimizer = optim.Adam( model.parameters(), lr=args.lr, betas=(0.9, 0.999), ) optimizer.zero_grad() #optimizer_D = optim.Adam( # model_D.parameters(), # lr=args.lr_D, # betas=(0.9, 0.999), # ) # optimizer_D.zero_grad() if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None print("===================================") print("====== Loading Testing Data =======") print("===================================") testset = ModelNetDatasetGT( root_list=args.test_file, sample_list=None, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) criterion = torch.nn.CrossEntropyLoss().to(device) train_pth = glob.glob(args.exp_dir + "/model_train_epoch_*.pth") print("Total #models: {}".format(len(train_pth))) max_accu = float("-inf") for i in range(len(train_pth)): print("Processing {} .....".format(i)) print("model {} .........".format(train_pth[i])) model.load_state_dict(torch.load(train_pth[i])) curr_accu, curr_loss = run_testing( dataloader=testloader, model=model, criterion=criterion, logger=test_logger, test_iter=100000000, writer=None, args=args, ) if curr_accu > max_accu: max_accu = curr_accu max_pth = train_pth[i] print("Max accuracy: {:.4f}".format(max_accu)) print("Trained model: {}".format(max_pth))
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device train_logger = make_logger("Train.log", args) test_logger = make_logger("Test.log", args) if args.train or args.train_stack: if args.train_regu: model = load_models( mode="seg_regu", device=device, args=args, ) else: model = load_models( mode="seg", device=device, args=args, ) model_D = load_models( mode="disc_stack", device=device, args=args, ) optimizer = optim.Adam( model.parameters(), lr=args.lr, betas=(0.9, 0.999), ) optimizer.zero_grad() optimizer_D = optim.Adam( model_D.parameters(), lr=args.lr_D_point, betas=(0.9, 0.999), ) optimizer_D.zero_grad() if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None if (args.train or args.train_stack) and args.test: print("===================================") print("====== Loading Training Data ======") print("===================================") if args.gt_sample_list != None: sample_gt_list = np.load(args.gt_sample_list) else: sample_gt_list = None trainset_gt = ShapeNetDatasetGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainset_nogt = ShapeNetDataset_noGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainloader_gt = torch.utils.data.DataLoader( trainset_gt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) trainloader_nogt = torch.utils.data.DataLoader( trainset_nogt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) print("===================================") print("====== Loading Test Data ======") print("===================================") testset = ShapeNetDatasetGT( root_list=args.test_file, sample_list=None, num_classes=16, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) args.iter_per_epoch = int(1.0 * trainset_gt.__len__() / args.batch_size) args.total_data = trainset_gt.__len__() + trainset_nogt.__len__() args.total_iterations = int(args.num_epochs * args.total_data / args.batch_size) args.iter_save_epoch = args.save_per_epoch * int(args.total_data / args.batch_size) args.iter_test_epoch = args.test_epoch * int(args.total_data / args.batch_size) args.semi_start = int(args.semi_start_epoch * args.total_data / args.batch_size) if (args.train or args.train_stack) and args.test: model.train() model_D.train() model.to(args.device) model_D.to(args.device) #class_weight = 1.0 / trainset_gt.get_class_probability().to(device) seg_loss = torch.nn.CrossEntropyLoss().to(device) gan_loss = torch.nn.BCEWithLogitsLoss().to(device) semi_loss = torch.nn.CrossEntropyLoss(ignore_index=255) history_pool_gt = ImagePool(args.pool_size) history_pool_nogt = ImagePool(args.pool_size) trainloader_gt_iter = enumerate(trainloader_gt) targetloader_nogt_iter = enumerate(trainloader_nogt) if args.train and (args.semi_start_epoch==0): run_training_seg( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, testdataset=testset, model=model, model_D=model_D, gan_loss=gan_loss, seg_loss=seg_loss, optimizer=optimizer, optimizer_D=optimizer_D, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) elif args.train and (args.semi_start_epoch>0): run_training_seg_semi( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, testdataset=testset, model=model, model_D=model_D, gan_loss=gan_loss, seg_loss=seg_loss, semi_loss=semi_loss, optimizer=optimizer, optimizer_D=optimizer_D, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) elif args.train_stack: gan_loss = torch.nn.BCELoss().to(device) shape_criterion = torch.nn.CrossEntropyLoss().to(device) regu_loss = torch.nn.MSELoss().to(device) if args.semi_start_epoch == 0: if args.train_regu: run_training_seg_stack_regulization( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, testdataset=testset, model=model, model_D=model_D, gan_loss=gan_loss, seg_loss=seg_loss, regu_loss=regu_loss, shape_criterion=shape_criterion, optimizer=optimizer, optimizer_D=optimizer_D, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) else: run_training_seg_stack( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, testdataset=testset, model=model, model_D=model_D, gan_loss=gan_loss, seg_loss=seg_loss, shape_criterion=shape_criterion, optimizer=optimizer, optimizer_D=optimizer_D, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) else: if args.train_regu: run_training_seg_stack_regulization_semi( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, testdataset=testset, model=model, model_D=model_D, gan_loss=gan_loss, seg_loss=seg_loss, regu_loss=regu_loss, semi_loss=semi_loss, shape_criterion=shape_criterion, optimizer=optimizer, optimizer_D=optimizer_D, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) else: run_training_seg_stack_semi( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, testdataset=testset, model=model, model_D=model_D, gan_loss=gan_loss, seg_loss=seg_loss, semi_loss=semi_loss, shape_criterion=shape_criterion, optimizer=optimizer, optimizer_D=optimizer_D, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) if args.dual: model = load_models( mode="seg", device=device, args=args, ) optimizer = optim.Adam( model.parameters(), lr=args.lr, betas=(0.9, 0.999), ) optimizer.zero_grad() sharedDisc, shapeDisc, pointDisc = load_models( mode="disc_dual", device=device, args=args, ) D_params_shape = (list(shapeDisc.parameters()) + list(sharedDisc.parameters())) D_params_point = (list(pointDisc.parameters()) + list(sharedDisc.parameters())) optimizer_D_shape = optim.SGD( D_params_shape, lr=args.lr_D_shape, # lr=args.lr_D*0.2, ) optimizer_D_shape.zero_grad() # optimizer_D_point = optim.Adam( # D_params_point, # lr=args.lr_D, # betas=(0.9, 0.999), # ) optimizer_D_point = optim.SGD( D_params_point, lr=args.lr_D_point, ) optimizer_D_point.zero_grad() print("===================================") print("====== Loading Training Data ======") print("===================================") if args.gt_sample_list != None: sample_gt_list = np.load(args.gt_sample_list) else: sample_gt_list = None trainset_gt = ShapeNetDatasetGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainset_nogt = ShapeNetDataset_noGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainloader_gt = torch.utils.data.DataLoader( trainset_gt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) trainloader_nogt = torch.utils.data.DataLoader( trainset_nogt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) print("===================================") print("====== Loading Test Data ======") print("===================================") testset = ShapeNetDatasetGT( root_list=args.test_file, sample_list=None, num_classes=16, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) args.iter_per_epoch = int(1.0 * trainset_gt.__len__() / args.batch_size) args.total_data = trainset_gt.__len__() + trainset_nogt.__len__() args.total_iterations = int(args.num_epochs * args.total_data / args.batch_size) args.iter_save_epoch = args.save_per_epoch * int(args.total_data / args.batch_size) args.iter_test_epoch = args.test_epoch * int(args.total_data / args.batch_size) args.semi_start = int(args.semi_start_epoch * args.total_data / args.batch_size) model.train() sharedDisc.train() shapeDisc.train() pointDisc.train() model.to(args.device) sharedDisc.to(args.device) shapeDisc.to(args.device) pointDisc.to(args.device) seg_loss = torch.nn.CrossEntropyLoss().to(device) gan_point_loss = torch.nn.BCEWithLogitsLoss().to(device) gan_shape_loss = torch.nn.CrossEntropyLoss().to(device) history_pool_gt = ImagePool(args.pool_size) history_pool_nogt = ImagePool(args.pool_size) trainloader_gt_iter = enumerate(trainloader_gt) targetloader_nogt_iter = enumerate(trainloader_nogt) run_training_seg_dual( trainloader_gt=trainloader_gt, trainloader_nogt=trainloader_nogt, trainloader_gt_iter=trainloader_gt_iter, targetloader_nogt_iter=targetloader_nogt_iter, testloader=testloader, testdataset=testset, model=model, sharedDisc=sharedDisc, shapeDisc=shapeDisc, pointDisc=pointDisc, gan_point_loss=gan_point_loss, gan_shape_loss=gan_shape_loss, seg_loss=seg_loss, optimizer=optimizer, optimizer_D_shape=optimizer_D_shape, optimizer_D_point=optimizer_D_point, history_pool_gt=history_pool_gt, history_pool_nogt=history_pool_nogt, writer=writer, train_logger=train_logger, test_logger=test_logger, args=args, ) if args.test: print("===================================") print("====== Loading Testing Data =======") print("===================================") testset = ShapeNetDatasetGT( root_list=args.test_file, sample_list=None, num_classes=16, ) testloader = torch.utils.data.DataLoader( testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, ) criterion = torch.nn.CrossEntropyLoss().to(device) run_testing_seg( dataset=testset, dataloader=testloader, model=model, criterion=criterion, logger=test_logger, test_iter=100000000, writer=None, args=args, ) if args.tsne: from utils.metric import batch_get_iou, object_names from torch.autograd import Variable model = load_models( mode="seg", device=device, args=args, ) args.batch_size = 1 labels = [] objects = [] if args.gt_sample_list != None: sample_gt_list = np.load(args.gt_sample_list) else: sample_gt_list = None trainset_gt = ShapeNetDatasetGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainset_nogt = ShapeNetDataset_noGT( root_list=args.train_file, sample_list=sample_gt_list, num_classes=16, ) trainloader_gt = torch.utils.data.DataLoader( trainset_gt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) trainloader_nogt = torch.utils.data.DataLoader( trainset_nogt, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) model.eval() shape_ious = np.empty(len(object_names), dtype=np.object) for i in range(shape_ious.shape[0]): shape_ious[i] = [] all_shapes_train_gt = np.empty((len(trainset_gt), 2048)) for batch_idx, data in enumerate(trainloader_gt): if batch_idx % 1000==0: print("Processing {} ...".format(batch_idx)) pts, cls, seg = data pts, cls, seg = Variable(pts).float(), \ Variable(cls), Variable(seg).type(torch.LongTensor) pts, cls, seg = pts.to(args.device), cls.to(args.device), seg.long().to(args.device) labels.append(1) objects.append(int(cls.argmax(axis=2).squeeze().cpu().numpy())) with torch.set_grad_enabled(False): pred, global_shape = model(pts, cls) all_shapes_train_gt[batch_idx,:] = global_shape.squeeze().detach().cpu().numpy() all_shapes_train_nogt = np.empty((len(trainset_nogt), 2048)) for batch_idx, data in enumerate(trainloader_nogt): if batch_idx % 1000==0: print("Processing {} ...".format(batch_idx)) pts, cls = data pts, cls = Variable(pts).float(), Variable(cls) pts, cls = pts.to(args.device), cls.to(args.device) with torch.set_grad_enabled(False): pred, global_shape = model(pts, cls) all_shapes_train_nogt[batch_idx, :] = global_shape.squeeze().detach().cpu().numpy() labels.append(0) objects.append(int(cls.argmax(axis=2).squeeze().cpu().numpy())) all_shapes = np.concatenate((all_shapes_train_gt, all_shapes_train_nogt), axis=0) shape_info = {"shapes": all_shapes, "labels": labels, "objects": objects} import pickle try: o = open("stack_global_shape_{}.pkl".format(len(trainset_gt)), "wb") pickle.dump(shape_info,o,protocol=2) o.close() except FileNotFoundError as e: print (e)