def run(config):
model = get_model(config).cuda()
criterion = get_loss(config)
optimizer = get_optimizer(config, model.parameters())
checkpoint = utils.checkpoint.get_initial_checkpoint(config)
if checkpoint is not None:
last_epoch, score = utils.checkpoint.load_checkpoint(config, model, checkpoint)
else:
print('[*] no checkpoint found')
last_epoch, score = -1, -1
print('last epoch:{} score:{:.4f}'.format(last_epoch, score))
optimizer.param_groups[0]['initial_lr'] = config.OPTIMIZER.LR
scheduler = get_scheduler(config, optimizer, last_epoch)
if last_epoch != -1:
scheduler.step()
writer = SummaryWriter(os.path.join(config.TRAIN_DIR, 'logs'))
train_loader = get_dataloader(config, 'train', transform=transforms.Compose([Albu(),
Normalize(),
ToTensor()]))
test_loader = get_dataloader(config, 'val', transform=transforms.Compose([Normalize(),
ToTensor()]))
train(config, model, train_loader, test_loader, criterion, optimizer, scheduler, writer, last_epoch+1, score)
def run(config):
model = get_model(config).to(device)
criterion = get_loss(config.LOSS.NAME)
optimizer = get_optimizer(config, model.parameters())
checkpoint = utils.checkpoint.get_initial_checkpoint(config)
if checkpoint is not None:
last_epoch, score, loss = utils.checkpoint.load_checkpoint(
config, model, checkpoint)
else:
print('[*] no checkpoint found')
last_epoch, score, loss = -1, -1, float('inf')
print('last epoch:{} score:{:.4f} loss:{:.4f}'.format(
last_epoch, score, loss))
optimizer.param_groups[0]['initial_lr'] = config.OPTIMIZER.LR
scheduler = get_scheduler(config, optimizer, last_epoch)
if config.SCHEDULER.NAME == 'multi_step':
milestones = scheduler.state_dict()['milestones']
step_count = len([i for i in milestones if i < last_epoch])
optimizer.param_groups[0]['lr'] *= scheduler.state_dict(
)['gamma']**step_count
if last_epoch != -1:
scheduler.step()
writer = SummaryWriter(os.path.join(config.TRAIN_DIR, 'logs'))
train_loader = get_dataloader(config,
'train',
transform=transforms.Compose([
Albu(),
CV2_Resize(config.DATA.IMG_W,
config.DATA.IMG_H),
Normalize(),
ToTensor()
]))
val_loader = get_dataloader(config,
'val',
transform=transforms.Compose([
CV2_Resize(config.DATA.IMG_W,
config.DATA.IMG_H),
Normalize(),
ToTensor()
]))
train(config, model, train_loader, val_loader, criterion, optimizer,
scheduler, writer, last_epoch + 1, score, loss)
def process_split(split='test', descriptor='fcgf'):
if split == 'test':
dset = ThreeDMatchTest(root='/data/3DMatch',
descriptor=descriptor,
num_node='all',
augment_axis=0,
augment_rotation=0.0,
augment_translation=0.0,
)
else:
dset = ThreeDMatchTrainVal(root='/data/3DMatch',
descriptor=descriptor,
split=split,
num_node='all',
augment_axis=0,
augment_rotation=0.0,
augment_translation=0.0,
)
dloader = get_dataloader(dset, batch_size=1, num_workers=16, shuffle=False)
dloader_iter = dloader.__iter__()
inlier_ratio_list = []
for i in tqdm(range(len(dset))):
corr, src_keypts, tgt_keypts, gt_trans, gt_labels = dloader_iter.next()
inlier_ratio = gt_labels.mean()
inlier_ratio_list.append(float(inlier_ratio)*100)
return np.array(inlier_ratio_list)
def eval_KITTI(method, args):
dset = KITTIDataset(
root=kitti_dir,
split='test',
descriptor=args.descriptor,
in_dim=6,
inlier_threshold=args.inlier_threshold,
num_node=15000,
use_mutual=args.use_mutual,
augment_axis=0,
augment_rotation=0.00,
augment_translation=0.0,
)
dloader = get_dataloader(dset, batch_size=1, num_workers=8, shuffle=False)
stats = eval_KITTI_scene(method, dloader, args)
# pair level average
allpair_stats = stats
allpair_average = allpair_stats.mean(0)
correct_pair_average = allpair_stats[allpair_stats[:, 0] == 1].mean(0)
logging.info(f"*" * 40)
logging.info(
f"All {allpair_stats.shape[0]} pairs, Mean Reg Recall={allpair_average[0] * 100:.2f}%, Mean Re={correct_pair_average[1]:.2f}, Mean Te={correct_pair_average[2]:.2f}"
)
logging.info(
f"\tInput: Mean Inlier Num={allpair_average[3]:.2f}(ratio={allpair_average[4] * 100:.2f}%)"
)
logging.info(
f"\tOutput: Mean Inlier Num={allpair_average[5]:.2f}(precision={allpair_average[6] * 100:.2f}%, recall={allpair_average[7] * 100:.2f}%, f1={allpair_average[8] * 100:.2f}%)"
)
logging.info(
f"\tMean model time: {allpair_average[9]:.2f}s, Mean data time: {allpair_average[10]:.2f}s"
)
return allpair_stats
def eval_KITTI(model, config, use_icp):
dset = KITTIDataset(root='/data/KITTI',
split='test',
descriptor=config.descriptor,
in_dim=config.in_dim,
inlier_threshold=config.inlier_threshold,
num_node=12000,
use_mutual=config.use_mutual,
augment_axis=0,
augment_rotation=0.00,
augment_translation=0.0,
)
dloader = get_dataloader(dset, batch_size=1, num_workers=16, shuffle=False)
stats = eval_KITTI_per_pair(model, dloader, config, use_icp)
logging.info(f"Max memory allicated: {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f}GB")
# pair level average
allpair_stats = stats
allpair_average = allpair_stats.mean(0)
correct_pair_average = allpair_stats[allpair_stats[:, 0] == 1].mean(0)
logging.info(f"*"*40)
logging.info(f"All {allpair_stats.shape[0]} pairs, Mean Success Rate={allpair_average[0]*100:.2f}%, Mean Re={correct_pair_average[1]:.2f}, Mean Te={correct_pair_average[2]:.2f}")
logging.info(f"\tInput: Mean Inlier Num={allpair_average[3]:.2f}(ratio={allpair_average[4]*100:.2f}%)")
logging.info(f"\tOutput: Mean Inlier Num={allpair_average[5]:.2f}(precision={allpair_average[6]*100:.2f}%, recall={allpair_average[7]*100:.2f}%, f1={allpair_average[8]*100:.2f}%)")
logging.info(f"\tMean model time: {allpair_average[9]:.2f}s, Mean data time: {allpair_average[10]:.2f}s")
return allpair_stats
def run(config):
model = get_model(config).to(device)
# model_params = [{'params': model.encoder.parameters(), 'lr': config.OPTIMIZER.ENCODER_LR},
# {'params': model.decoder.parameters(), 'lr': config.OPTIMIZER.DECODER_LR}]
optimizer = get_optimizer(config, model.parameters())
# optimizer = get_optimizer(config, model_params)
checkpoint = utils.checkpoint.get_initial_checkpoint(config)
if checkpoint is not None:
last_epoch, score, loss = utils.checkpoint.load_checkpoint(config, model, checkpoint)
else:
print('[*] no checkpoint found')
last_epoch, score, loss = -1, -1, float('inf')
print('last epoch:{} score:{:.4f} loss:{:.4f}'.format(last_epoch, score, loss))
optimizer.param_groups[0]['initial_lr'] = config.OPTIMIZER.LR
# optimizer.param_groups[0]['initial_lr'] = config.OPTIMIZER.ENCODER_LR
# optimizer.param_groups[1]['initial_lr'] = config.OPTIMIZER.DECODER_LR
scheduler = get_scheduler(config, optimizer, last_epoch)
if config.SCHEDULER.NAME == 'multi_step':
milestones = scheduler.state_dict()['milestones']
step_count = len([i for i in milestones if i < last_epoch])
optimizer.param_groups[0]['lr'] *= scheduler.state_dict()['gamma'] ** step_count
# optimizer.param_groups[0]['lr'] *= scheduler.state_dict()['gamma'] ** step_count
# optimizer.param_groups[1]['lr'] *= scheduler.state_dict()['gamma'] ** step_count
if last_epoch != -1:
scheduler.step()
log_train = Logger()
log_val = Logger()
log_train.open(os.path.join(config.TRAIN_DIR, 'log_train.txt'), mode='a')
log_val.open(os.path.join(config.TRAIN_DIR, 'log_val.txt'), mode='a')
train_loader = get_dataloader(config, 'train', transform=Albu(config.ALBU))
val_loader = get_dataloader(config, 'val')
train(config, model, train_loader, val_loader, optimizer, scheduler, log_train, log_val, last_epoch+1, score, loss)
def run(config):
model = get_model(config).to(device)
optimizer = get_optimizer(config, model.parameters())
checkpoint = utils.checkpoint.get_initial_checkpoint(config)
if checkpoint is not None:
last_epoch, score, loss = utils.checkpoint.load_checkpoint(
config, model, checkpoint)
else:
print('[*] no checkpoint found')
last_epoch, score, loss = -1, -1, float('inf')
print('last epoch:{} score:{:.4f} loss:{:.4f}'.format(
last_epoch, score, loss))
optimizer.param_groups[0]['initial_lr'] = config.OPTIMIZER.LR
scheduler = get_scheduler(config, optimizer, last_epoch)
if config.SCHEDULER.NAME == 'multi_step':
milestones = scheduler.state_dict()['milestones']
step_count = len([i for i in milestones if i < last_epoch])
optimizer.param_groups[0]['lr'] *= scheduler.state_dict(
)['gamma']**step_count
if last_epoch != -1:
scheduler.step()
# writer = SummaryWriter(os.path.join(config.TRAIN_DIR, 'logs'))
log_train = Logger()
log_val = Logger()
log_train.open(os.path.join(config.TRAIN_DIR, 'log_train.txt'), mode='a')
log_val.open(os.path.join(config.TRAIN_DIR, 'log_val.txt'), mode='a')
augmentation = Albu_Seg() if config.TASK == 'seg' else Albu_Cls()
train_loader = get_dataloader(config, 'train', transform=augmentation)
val_loader = get_dataloader(config, 'val')
train(config, model, train_loader, val_loader, optimizer, scheduler,
log_train, log_val, last_epoch + 1, score, loss)
def run(config, num_checkpoint, epoch_end, output_filename):
dataloader = get_dataloader(config, split='val', transform=None)
model = get_model(config).cuda()
checkpoints = get_checkpoints(config, num_checkpoint, epoch_end)
utils.checkpoint.load_checkpoint(config, model, checkpoints[0])
for i, checkpoint in enumerate(checkpoints[1:]):
model2 = get_model(config).cuda()
last_epoch, _, _ = utils.checkpoint.load_checkpoint(config, model2, checkpoint)
swa.moving_average(model, model2, 1. / (i + 2))
with torch.no_grad():
swa.bn_update(dataloader, model)
# output_name = '{}.{}.{:03d}'.format(output_filename, num_checkpoint, last_epoch)
# print('save {}'.format(output_name))
utils.checkpoint.save_checkpoint(config, model, None, None, epoch_end,
weights_dict={'state_dict': model.state_dict()},
name=output_filename)
def eval_3DMatch(model, config, args):
dset = ThreeDLOMatchTest(
root='/data/3DMatch',
descriptor='fcgf',
in_dim=config.in_dim,
inlier_threshold=config.inlier_threshold,
num_node=args.num_points,
use_mutual=config.use_mutual,
augment_axis=0,
augment_rotation=0.00,
augment_translation=0.0,
)
dloader = get_dataloader(dset, batch_size=1, num_workers=16, shuffle=False)
allpair_stats, allpair_poses = eval_3DMatch_scene(model, 0, dloader,
config, args)
# benchmarking using the registration recall defined in 3DMatch to compare with Predator
# np.save('predator.npy', allpair_poses)
benchmark_predator(
allpair_poses,
gt_folder='/data/OverlapPredator/configs/benchmarks/3DLoMatch')
# benchmarking using the registration recall defined in DGR
allpair_average = allpair_stats.mean(0)
correct_pair_average = allpair_stats[allpair_stats[:, 0] == 1].mean(0)
logging.info(f"*" * 40)
logging.info(
f"All {allpair_stats.shape[0]} pairs, Mean Success Rate={allpair_average[0] * 100:.2f}%, Mean Re={correct_pair_average[1]:.2f}, Mean Te={correct_pair_average[2]:.2f}"
)
logging.info(
f"\tInput: Mean Inlier Num={allpair_average[3]:.2f}(ratio={allpair_average[4] * 100:.2f}%)"
)
logging.info(
f"\tOutput: Mean Inlier Num={allpair_average[5]:.2f}(precision={allpair_average[6] * 100:.2f}%, recall={allpair_average[7] * 100:.2f}%, f1={allpair_average[8] * 100:.2f}%)"
)
logging.info(
f"\tMean model time: {allpair_average[9]:.2f}s, Mean data time: {allpair_average[10]:.2f}s"
)
return allpair_stats
if args.generate_features:
import importlib
cfg = importlib.import_module("configs.config")
dset = KITTIMapDataset("test",
cfg,
config_d3feat=config,
root=config.root)
# dset = ThreeDMatchTestset(root=config.root,
# downsample=config.downsample,
# config=config,
# last_scene=False,
# )
dloader, _ = get_dataloader(
dataset=dset,
batch_size=1,
shuffle=False,
num_workers=config.num_workers,
)
generate_features(model.cuda(), dloader, config, args.chosen_snapshot)
def test_kitti(model, dataset, config):
# self.sess.run(dataset.test_init_op)
import sys
use_random_points = False
icp_save_path = "d3feat_output"
if use_random_points:
num_keypts = 5000
# icp_save_path = f'geometric_registration_kitti/D3Feat_{self.experiment_str}-rand{num_keypts}'
else:
num_keypts = 250
# icp_save_path = f'geometric_registration_kitti/D3Feat_{self.experiment_str}-pred{num_keypts}'
def main(args):
args.best_acc = 0
best_acc5 = 0
# Data
print('==> Preparing data..')
train_loader, test_loader, ndata = get_dataloader(args, add_erasing=args.erasing, aug_plus=args.aug_plus)
logger.info(f"length of training dataset: {ndata}")
# Model
model, model_ema = build_model(args)
contrast = MemoryMoCo(128, args.nce_k, args.nce_t, thresh=0).cuda()
criterion = NCESoftmaxLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
lr=args.batch_size * dist.get_world_size() / 128 * args.base_learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = get_scheduler(optimizer, len(train_loader), args)
if args.amp_opt_level != "O0":
if amp is None:
logger.warning(f"apex is not installed but amp_opt_level is set to {args.amp_opt_level}, ignoring.\n"
"you should install apex from https://github.com/NVIDIA/apex#quick-start first")
args.amp_opt_level = "O0"
else:
model, optimizer = amp.initialize(model, optimizer, opt_level=args.amp_opt_level)
model_ema = amp.initialize(model_ema, opt_level=args.amp_opt_level)
model = DistributedDataParallel(model, device_ids=[args.local_rank], broadcast_buffers=False)
# optionally resume from a checkpoint
if args.resume:
assert os.path.isfile(args.resume)
load_checkpoint(args, model, model_ema, contrast, optimizer, scheduler)
# tensorboard
if dist.get_rank() == 0:
summary_writer = SummaryWriter(log_dir=args.save_dir)
else:
summary_writer = None
# routine
for epoch in range(args.start_epoch, args.epochs + 1):
if args.lr_scheduler == 'cosine':
train_loader.sampler.set_epoch(epoch)
tic = time.time()
loss, prob = train_moco(epoch, train_loader, model, model_ema, contrast, criterion, optimizer, scheduler, args)
logger.info('epoch {}, total time {:.2f}'.format(epoch, time.time() - tic))
if summary_writer is not None:
# tensorboard logger
summary_writer.add_scalar('ins_loss', loss, epoch)
summary_writer.add_scalar('ins_prob', prob, epoch)
summary_writer.add_scalar('learning_rate', optimizer.param_groups[0]['lr'], epoch)
if args.dataset == 'stl10-full':
acc, acc5 = kNN(epoch, model, contrast, labeledTrainloader, test_loader, 200, args.nce_t, True)
else:
acc, acc5 = kNN(epoch, model, contrast, train_loader, test_loader, 200, args.nce_t, True)
if acc >= args.best_acc:
args.best_acc = acc
best_acc5 = acc5
logger.info('KNN top-1 precion: {:.4f} {:.4f}, best is: {:.4f} {:.4f}'.format(acc*100., \
acc5*100., args.best_acc*100., best_acc5*100))
logger.info(str(args))
if dist.get_rank() == 0:
# save model
save_checkpoint(args, epoch, model, model_ema, contrast, optimizer, scheduler, args.best_acc)
if args.dataset == 'stl10-full':
acc1, acc5 = kNN(epoch, model, contrast, labeledTrainloader, test_loader, 200, args.nce_t, True)
else:
acc1, acc5 = kNN(epoch, model, contrast, train_loader, test_loader, 200, args.nce_t, True)
logger.info('KNN top-1 and top-5 precion with recomputed memory bank: {:.4f} {:.4f}'.format(acc1*100., acc5*100))
logger.info('Best KNN top-1 and top-5 precion: {:.4f} {:.4f}'.format(args.best_acc*100., best_acc5*100))
logger.info(str(args))
config.optimizer,
gamma=config.scheduler_gamma,
)
# create dataset and dataloader
info_train = load_obj(config.train_info)
info_val = load_obj(config.val_info)
info_benchmark = load_obj(f'configs/{config.test_info}.pkl')
train_set = ThreeDMatchDownsampled(info_train,config,data_augmentation=True)
val_set = ThreeDMatchDownsampled(info_val,config,data_augmentation=False)
benchmark_set = ThreeDMatchDownsampled(info_benchmark,config, data_augmentation=False)
config.train_loader, neighborhood_limits = get_dataloader(dataset=train_set,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_workers,
)
config.val_loader, _ = get_dataloader(dataset=val_set,
batch_size=config.batch_size,
shuffle=False,
num_workers=1,
neighborhood_limits=neighborhood_limits
)
config.test_loader, _ = get_dataloader(dataset=benchmark_set,
batch_size=config.batch_size,
shuffle=False,
num_workers=1,
neighborhood_limits=neighborhood_limits)
# create evaluation metrics
config.architecture.append('resnetb')
for i in range(config.num_layers - 2):
config.architecture.append('nearest_upsample')
config.architecture.append('unary')
config.architecture.append('nearest_upsample')
config.architecture.append('last_unary')
config.model = KPFCNN(config).to(config.device)
# create dataset and dataloader
info_train = load_obj(config.train_info)
train_set = IndoorDataset(info_train, config, data_augmentation=True)
demo_set = ThreeDMatchDemo(config, config.src_pcd, config.tgt_pcd)
_, neighborhood_limits = get_dataloader(
dataset=train_set,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_workers,
)
demo_loader, _ = get_dataloader(dataset=demo_set,
batch_size=config.batch_size,
shuffle=False,
num_workers=1,
neighborhood_limits=neighborhood_limits)
# load pretrained weights
assert config.pretrain != None
state = torch.load(config.pretrain)
config.model.load_state_dict(state['state_dict'])
# do pose estimation
main(config, demo_loader)
def eval_3DMatch(method, args):
"""
Collect the evaluation results on each scene of 3DMatch testset, write the result to a .log file.
"""
scene_list = [
'7-scenes-redkitchen', 'sun3d-home_at-home_at_scan1_2013_jan_1',
'sun3d-home_md-home_md_scan9_2012_sep_30', 'sun3d-hotel_uc-scan3',
'sun3d-hotel_umd-maryland_hotel1', 'sun3d-hotel_umd-maryland_hotel3',
'sun3d-mit_76_studyroom-76-1studyroom2',
'sun3d-mit_lab_hj-lab_hj_tea_nov_2_2012_scan1_erika'
]
all_stats = {}
for scene_ind, scene in enumerate(scene_list):
dset = ThreeDMatchTest(
root='/data/3DMatch/',
descriptor=args.descriptor,
in_dim=6,
inlier_threshold=args.inlier_threshold,
num_node='all',
use_mutual=args.use_mutual,
augment_axis=0,
augment_rotation=0.00,
augment_translation=0.0,
select_scene=scene,
)
dloader = get_dataloader(dset,
batch_size=1,
num_workers=8,
shuffle=False)
scene_stats = eval_3DMatch_scene(method, scene, scene_ind, dloader,
args)
all_stats[scene] = scene_stats
scene_vals = np.zeros([len(scene_list), 12])
scene_ind = 0
for scene, stats in all_stats.items():
correct_pair = np.where(stats[:, 0] == 1)
scene_vals[scene_ind] = stats.mean(0)
# for Re and Te, we only count the matched pairs.
scene_vals[scene_ind, 1] = stats[correct_pair].mean(0)[1]
scene_vals[scene_ind, 2] = stats[correct_pair].mean(0)[2]
logging.info(f"Scene {scene_ind}th:"
f" Reg Recall={scene_vals[scene_ind, 0]*100:.2f}% "
f" Mean RE={scene_vals[scene_ind, 1]:.2f} "
f" Mean TE={scene_vals[scene_ind, 2]:.2f} "
f" Mean Precision={scene_vals[scene_ind, 6]*100:.2f}% "
f" Mean Recall={scene_vals[scene_ind, 7]*100:.2f}% "
f" Mean F1={scene_vals[scene_ind, 8]*100:.2f}%")
scene_ind += 1
# scene level average
average = scene_vals.mean(0)
logging.info(
f"All {len(scene_list)} scenes, Mean Reg Recall={average[0]*100:.2f}%, Mean Re={average[1]:.2f}, Mean Te={average[2]:.2f}"
)
logging.info(
f"\tInput: Mean Inlier Num={average[3]:.2f}(ratio={average[4]*100:.2f}%)"
)
logging.info(
f"\tOutput: Mean Inlier Num={average[5]:.2f}(precision={average[6]*100:.2f}%, recall={average[7]*100:.2f}%, f1={average[8]*100:.2f}%)"
)
logging.info(
f"\tMean model time: {average[9]:.4f}s, Mean data time: {average[10]:.4f}s"
)
# pair level average
stats_list = [stats for _, stats in all_stats.items()]
allpair_stats = np.concatenate(stats_list, axis=0)
allpair_average = allpair_stats.mean(0)
correct_pair_average = allpair_stats[allpair_stats[:, 0] == 1].mean(0)
logging.info(f"*" * 40)
logging.info(
f"All {allpair_stats.shape[0]} pairs, Mean Reg Recall={allpair_average[0]*100:.2f}%, Mean Re={correct_pair_average[1]:.2f}, Mean Te={correct_pair_average[2]:.2f}"
)
logging.info(
f"\tInput: Mean Inlier Num={allpair_average[3]:.2f}(ratio={allpair_average[4]*100:.2f}%)"
)
logging.info(
f"\tOutput: Mean Inlier Num={allpair_average[5]:.2f}(precision={allpair_average[6]*100:.2f}%, recall={allpair_average[7]*100:.2f}%, f1={allpair_average[8]*100:.2f}%)"
)
logging.info(
f"\tMean model time: {allpair_average[9]:.4f}s, Mean data time: {allpair_average[10]:.4f}s"
)
all_stats_npy = np.concatenate([v for k, v in all_stats.items()], axis=0)
return all_stats_npy
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
best_acc1 = 0 # best test accuracy
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
# Data
print('==> Preparing data..')
trainloader, testloader, ndata = get_dataloader(args)
print('==> Building model..')
net = models.__dict__['ResNet18'](low_dim=args.low_dim, pool_len=args.pool_len)
# define leminiscate
if args.nce_k > 0:
lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t,
args.nce_m)
else:
lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t, args.nce_m)
net.to(device)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
def __init__(self, config):
is_test = False
if is_test:
self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M') + 'Test'
else:
self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M')
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.verbose = True
# snapshot
self.snapshot_interval = 5
snapshot_root = f'snapshot/{config.dataset}_{self.experiment_id}'
tensorboard_root = f'tensorboard/{config.dataset}_{self.experiment_id}'
os.makedirs(snapshot_root, exist_ok=True)
os.makedirs(tensorboard_root, exist_ok=True)
shutil.copy2(os.path.join('.', 'training_ShapeNetCls.py'), os.path.join(snapshot_root, 'train.py'))
shutil.copy2(os.path.join('datasets', 'ShapeNet.py'), os.path.join(snapshot_root, 'dataset.py'))
shutil.copy2(os.path.join('datasets', 'dataloader.py'), os.path.join(snapshot_root, 'dataloader.py'))
self.save_dir = os.path.join(snapshot_root, 'models/')
self.result_dir = os.path.join(snapshot_root, 'results/')
self.tboard_dir = tensorboard_root
# dataset & dataloader
self.train_set = ShapeNetDataset(root=config.data_train_dir,
split='train',
first_subsampling_dl=config.first_subsampling_dl,
classification=True,
config=config,
)
self.test_set = ShapeNetDataset(root=config.data_test_dir,
split='test',
first_subsampling_dl=config.first_subsampling_dl,
classification=True,
config=config,
)
self.train_loader = get_dataloader(dataset=self.train_set,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=4,
)
self.test_loader = get_dataloader(dataset=self.test_set,
batch_size=config.test_batch_size,
shuffle=False,
num_workers=4,
)
print("Training set size:", self.train_loader.dataset.__len__())
print("Test set size:", self.test_loader.dataset.__len__())
# model
self.model = KPCNN(config)
self.resume = config.resume
# optimizer
self.start_epoch = 0
self.epoch = config.max_epoch
self.optimizer = optim.SGD(self.model.parameters(), lr=config.learning_rate, momentum=config.momentum, weight_decay=1e-6)
self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=config.exp_gamma)
self.scheduler_interval = config.exp_interval
# evaluate
self.evaluate_interval = 1
self.evaluate_metric = nn.CrossEntropyLoss(reduction='mean')
self.check_args()
augment_translation=config.augment_translation,
)
val_set = ThreeDMatchTrainVal(root=config.root,
split='val',
descriptor=config.descriptor,
in_dim=config.in_dim,
inlier_threshold=config.inlier_threshold,
num_node=config.num_node,
use_mutual=config.use_mutual,
downsample=config.downsample,
augment_axis=config.augment_axis,
augment_rotation=config.augment_rotation,
augment_translation=config.augment_translation,
)
config.train_loader = get_dataloader(dataset=train_set,
batch_size=config.batch_size,
num_workers=config.num_workers,
)
config.val_loader = get_dataloader(dataset=val_set,
batch_size=config.batch_size,
num_workers=config.num_workers,
)
# create evaluation
config.evaluate_metric = {
"ClassificationLoss": ClassificationLoss(balanced=config.balanced),
"SpectralMatchingLoss": SpectralMatchingLoss(balanced=config.balanced),
"TransformationLoss": TransformationLoss(re_thre=config.re_thre, te_thre=config.te_thre),
}
config.metric_weight = {
"ClassificationLoss": config.weight_classification,
"SpectralMatchingLoss": config.weight_spectralmatching,
def main():
seed_everything()
# yml = 'configs/base.yml'
# config = utils.config.load(yml)
# pprint.pprint(config, indent=2)
model = get_model(config).cuda()
bind_model(model)
args = get_args()
if args.pause: ## test mode일 때
print('Inferring Start...')
nsml.paused(scope=locals())
if args.mode == 'train': ### training mode일 때
print('Training Start...')
# no bias decay
if config.OPTIMIZER.NO_BIAS_DECAY:
group_decay, group_no_decay = group_weight(model)
params = [{'params': group_decay},
{'params': group_no_decay, 'weight_decay': 0.0}]
else:
params = model.parameters()
optimizer = get_optimizer(config, params)
optimizer.param_groups[0]['initial_lr'] = config.OPTIMIZER.LR
if config.OPTIMIZER.NO_BIAS_DECAY:
optimizer.param_groups[1]['initial_lr'] = config.OPTIMIZER.LR
###############################################################################################
if IS_LOCAL:
prepare_train_directories(config)
utils.config.save_config(yml, config.LOCAL_TRAIN_DIR)
checkpoint = utils.checkpoint.get_initial_checkpoint(config)
if checkpoint is not None:
last_epoch, score, loss = utils.checkpoint.load_checkpoint(config, model, checkpoint)
else:
print('[*] no checkpoint found')
last_epoch, score, loss = -1, -1, float('inf')
print('last epoch:{} score:{:.4f} loss:{:.4f}'.format(last_epoch, score, loss))
else:
last_epoch = -1
###############################################################################################
scheduler = get_scheduler(config, optimizer, last_epoch=last_epoch)
###############################################################################################
if IS_LOCAL:
if config.SCHEDULER.NAME == 'multi_step':
if config.SCHEDULER.WARMUP:
scheduler_dict = scheduler.state_dict()['after_scheduler'].state_dict()
else:
scheduler_dict = scheduler.state_dict()
milestones = scheduler_dict['milestones']
step_count = len([i for i in milestones if i < last_epoch])
optimizer.param_groups[0]['lr'] *= scheduler_dict['gamma'] ** step_count
if config.OPTIMIZER.NO_BIAS_DECAY:
optimizer.param_groups[1]['initial_lr'] *= scheduler_dict['gamma'] ** step_count
if last_epoch != -1:
scheduler.step()
###############################################################################################
# for dirname, _, filenames in os.walk(DATASET_PATH):
# for filename in filenames:
# print(os.path.join(dirname, filename))
# if preprocessing possible
preprocess_type = config.DATA.PREPROCESS
cv2_size = (config.DATA.IMG_W, config.DATA.IMG_H)
if not IS_LOCAL:
preprocess(os.path.join(DATASET_PATH, 'train', 'train_data', 'NOR'), os.path.join(preprocess_type, 'NOR'), preprocess_type, cv2_size)
preprocess(os.path.join(DATASET_PATH, 'train', 'train_data', 'AMD'), os.path.join(preprocess_type, 'AMD'), preprocess_type, cv2_size)
preprocess(os.path.join(DATASET_PATH, 'train', 'train_data', 'RVO'), os.path.join(preprocess_type, 'RVO'), preprocess_type, cv2_size)
preprocess(os.path.join(DATASET_PATH, 'train', 'train_data', 'DMR'), os.path.join(preprocess_type, 'DMR'), preprocess_type, cv2_size)
data_dir = preprocess_type
# data_dir = os.path.join(DATASET_PATH, 'train/train_data')
else: # IS_LOCAL
data_dir = os.path.join(DATASET_PATH, preprocess_type)
# eda
# train_std(data_dir, preprocess_type, cv2_size)
fold_df = split_cv(data_dir, n_splits=config.NUM_FOLDS)
val_fold_idx = config.IDX_FOLD
###############################################################################################
train_loader = get_dataloader(config, data_dir, fold_df, val_fold_idx, 'train', transform=Albu())
val_loader = get_dataloader(config, data_dir, fold_df, val_fold_idx, 'val')
postfix = dict()
num_epochs = config.TRAIN.NUM_EPOCHS
val_acc_list = []
for epoch in range(last_epoch+1, num_epochs):
if epoch >= config.LOSS.FINETUNE_EPOCH:
criterion = get_loss(config.LOSS.FINETUNE_LOSS)
else:
criterion = get_loss(config.LOSS.NAME)
train_values = train_single_epoch(config, model, train_loader, criterion, optimizer, scheduler, epoch)
val_values = evaluate_single_epoch(config, model, val_loader, criterion, epoch)
val_acc_list.append((epoch, val_values[2]))
if config.SCHEDULER.NAME != 'one_cyle_lr':
scheduler.step()
if IS_LOCAL:
utils.checkpoint.save_checkpoint(config, model, epoch, val_values[1], val_values[0])
else:
postfix['train_loss'] = train_values[0]
postfix['train_res'] = train_values[1]
postfix['train_acc'] = train_values[2]
postfix['train_sens'] = train_values[3]
postfix['train_spec'] = train_values[4]
postfix['val_loss'] = val_values[0]
postfix['val_res'] = val_values[1]
postfix['val_acc'] = val_values[2]
postfix['val_sens'] = val_values[3]
postfix['val_spec'] = val_values[4]
nsml.report(**postfix, summary=True, step=epoch)
val_res = '%.10f' % val_values[1]
val_res = val_res.replace(".", "")
val_res = val_res[:4] + '.' + val_res[4:8] + '.' + val_res[8:]
save_name = 'epoch_%02d_score%s_loss%.4f.pth' % (epoch, val_res, val_values[0])
# nsml.save(save_name)
nsml.save(epoch)
for e, val_acc in val_acc_list:
print('%02d %s' % (e, val_acc))
# return dict_inputs
return point_set, features, cls
else:
# manually convert numpy array to Tensor.
# seg = torch.from_numpy(seg) - 1 # change to 0-based labels
# dict_inputs = segmentation_inputs(point_set, features, seg, self.config)
# return dict_inputs
seg = seg - 1
return point_set, features, seg
def __len__(self):
return len(self.datapath)
if __name__ == '__main__':
datapath = "./data/shapenetcore_partanno_segmentation_benchmark_v0"
from training_ShapeNetCls import ShapeNetPartConfig
config = ShapeNetPartConfig()
print("Segmentation task:")
dset = ShapeNetDataset(root=datapath, config=config, first_subsampling_dl=0.01, classification=False)
input = dset[0]
from datasets.dataloader import get_dataloader
dataloader = get_dataloader(dset, batch_size=2)
for iter, input in enumerate(dataloader):
print(input)
break
