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)
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(os.path.join(RES_DIR, args.name),
"evaluation_trinity_test")
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
transforms_shape_target = dual_transforms.Compose([
dual_transforms.CenterCrop((400, 400)),
dual_transforms.Scale(args.image_size[0]),
])
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,
)
assert args.checkpoint_seg is not None, "Need trained .pth!"
model_seg = load_models(
mode="segmentation",
device=device,
args=args,
)
print("Evaluating ...................")
miou_all, iou_all = run_testing(
dataset=testdata,
test_loader=test_loader,
model=model_seg,
args=args,
)
# print('Global Mean Accuracy: {:.3f}'.format(np.array(pm.GA).mean()))
# print('Mean IOU: {:.3f}'.format(np.array(pm.IOU).mean()))
# print('Mean Recall: {:.3f}'.format(np.array(pm.Recall).mean()))
# print('Mean Precision: {:.3f}'.format(np.array(pm.Precision).mean()))
# print('Mean F1: {:.3f}'.format(np.array(pm.F1).mean()))()
print('Mean IOU: {:.3f}'.format(miou_all.mean()))
print("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(),
))