it = -1 # for the initialize value of `LambdaLR` and `BNMomentumScheduler`
best_loss = 1e10
start_epoch = 1
# load status from checkpoint
if args.checkpoint is not None:
checkpoint_status = pt_utils.load_checkpoint(
model, optimizer, filename=args.checkpoint.split(".")[0])
if checkpoint_status is not None:
it, start_epoch, best_loss = checkpoint_status
lr_scheduler = lr_sched.LambdaLR(optimizer,
lr_lambda=lr_lbmd,
last_epoch=it)
bnm_scheduler = pt_utils.BNMomentumScheduler(model,
bn_lambda=bn_lbmd,
last_epoch=it)
model = nn.DataParallel(model)
it = max(it, 0) # for the initialize value of `trainer.train`
model_fn = model_fn_decorator(nn.CrossEntropyLoss())
if args.visdom:
viz = pt_utils.VisdomViz(port=args.visdom_port)
else:
viz = pt_utils.CmdLineViz()
viz.text(pprint.pformat(vars(args)))
if not osp.isdir("checkpoints"):
def train(cmd_args):
if os.path.exists(cmd_args.log_dir) is not True:
os.makedirs(cmd_args.log_dir)
checkpoints_dir = os.path.join(cmd_args.log_dir, "checkpoints")
if os.path.exists(checkpoints_dir) is not True:
os.mkdir(checkpoints_dir)
lr_clip = 1e-5
bnm_clip = 1e-2
trans = transforms.Compose([
d_utils.PointcloudToTensor(),
#d_utils.PointcloudRandomPermutation(),
d_utils.PointcloudJitter()
])
train_set = bhcp_dataloader.bhcp_dataloader(cmd_args.data_dir,
cmd_args.category,
transforms=trans,
is_pts_aligned=False)
train_loader = DataLoader(
train_set,
batch_size=cmd_args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=False,
)
model = Pointnet2StructurePointNet(
num_structure_points=cmd_args.num_structure_points,
input_channels=0,
use_xyz=True)
model.cuda()
optimizer = optim.Adam(model.parameters(),
lr=cmd_args.lr,
weight_decay=cmd_args.weight_decay)
lr_lbmd = lambda it: max(
cmd_args.lr_decay**
(int(it * cmd_args.batch_size / cmd_args.decay_step)),
lr_clip / cmd_args.lr,
)
bn_lbmd = lambda it: max(
cmd_args.bn_momentum * cmd_args.bnm_decay**
(int(it * cmd_args.batch_size / cmd_args.decay_step)),
bnm_clip,
)
iters = -1
criterions = {'ComputeLoss3d': ComputeLoss3d()}
# load status from checkpoint
log_open_mode = 'w'
start_epoch = 0
if cmd_args.checkpoint is not None:
fname = os.path.join(checkpoints_dir, cmd_args.checkpoint)
start_epoch, iters = checkpoint_util.load_checkpoint(
model_3d=model, optimizer=optimizer, filename=fname)
start_epoch += 1
log_open_mode = 'a'
log = LogUtils(os.path.join(cmd_args.log_dir, 'logfile'), log_open_mode)
log.write('train unsupervised structure points for bhcp\n')
log.write_args(cmd_args)
lr_scheduler = lr_sched.LambdaLR(optimizer,
lr_lambda=lr_lbmd,
last_epoch=iters)
bnm_scheduler = pt_utils.BNMomentumScheduler(model,
bn_lambda=bn_lbmd,
last_epoch=iters)
iters = max(iters, 0)
for epoch_i in range(start_epoch, cmd_args.max_epochs):
log.write('\nepoch: {0} ###########################'.format(epoch_i))
train_loss, train_cd_loss, train_consistent_loss, iters = train_one_epoch(
model,
optimizer,
train_loader,
lr_scheduler,
bnm_scheduler,
iters,
criterions,
num_of_trans=cmd_args.num_of_transform,
num_inpts=cmd_args.num_inpts)
log.write('\nave_train_loss:{0}, cd_loss:{1}, consis_loss:{2}'.format(
("%.8f" % train_loss), ("%.8f" % train_cd_loss),
("%.8f" % train_consistent_loss)))
if cmd_args.checkpoint_save_step != -1 and (
epoch_i + 1) % cmd_args.checkpoint_save_step is 0:
fname = os.path.join(checkpoints_dir,
'checkpoint_{}'.format(epoch_i))
checkpoint_util.save_checkpoint(filename=fname,
model_3d=model,
optimizer=optimizer,
iters=iters,
epoch=epoch_i)
fname = os.path.join(checkpoints_dir, 'model')
checkpoint_util.save_checkpoint(filename=fname,
model_3d=model,
optimizer=optimizer,
iters=iters,
epoch=epoch_i)
batch_size=args.batch_size,
pin_memory=True,
num_workers=2,
shuffle=True)
model = Pointnet(num_classes=13, input_channels=6, use_xyz=True)
model.cuda()
optimizer = optim.Adam(
model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
lr_lbmd = lambda it: max(args.lr_decay**(int(it * args.batch_size / args.decay_step)), lr_clip / args.lr)
bnm_lmbd = lambda it: max(args.bn_momentum * args.bn_decay**(int(it * args.batch_size / args.decay_step)), bnm_clip)
if args.checkpoint is None:
lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lbmd)
bnm_scheduler = pt_utils.BNMomentumScheduler(model, bnm_lmbd)
start_epoch = 1
best_prec = 0
best_loss = 1e10
else:
start_epoch, best_loss = pt_utils.load_checkpoint(
model, optimizer, filename=args.checkpoint.split(".")[0])
lr_scheduler = lr_sched.LambdaLR(
optimizer, lr_lbmd, last_epoch=start_epoch)
bnm_scheduler = pt_utils.BNMomentumScheduler(
model, bnm_lmbd, last_epoch=start_epoch)
model_fn = model_fn_decorator(nn.CrossEntropyLoss())
if args.visdom:
