def evaluate(args):
cache = _load(root)
norm = True if args.model_name == 'pointnet' else False
test_ds = PartNormalDataset(root, cache, npoints=2048, split='test')
testdataloader = DataLoader(test_ds,
batch_size=args.batch_size,
shuffle=False,
num_workers=int(args.workers))
log.info("The number of test data is:", len(test_ds))
log.info('Building Model', args.model_name)
num_classes = 16
num_part = 50
if args.model_name == 'pointnet2':
model = PointNet2PartSegMsg_one_hot(num_part)
else:
model = PointNetDenseCls(cat_num=num_classes, part_num=num_part)
torch.backends.cudnn.benchmark = True
model = torch.nn.DataParallel(model).cuda()
log.debug('Using gpu:', args.gpu)
if args.pretrain is None:
log.err('No pretrain model')
return
log.debug('Loading pretrain model...')
state_dict = torch.load(args.pretrain)
model.load_state_dict(state_dict)
log.info('Testing pretrain model...')
test_metrics, test_hist_acc, cat_mean_iou = test_partseg(
model.eval(),
testdataloader,
label_id_to_name,
args.model_name,
num_part,
)
log.info('test_hist_acc', len(test_hist_acc))
log.info(cat_mean_iou)
log.info('Test Accuracy', '%.5f' % test_metrics['accuracy'])
log.info('Class avg mIOU:', '%.5f' % test_metrics['class_avg_iou'])
log.info('Inctance avg mIOU:', '%.5f' % test_metrics['inctance_avg_iou'])
def main(args):
os.environ[
"CUDA_VISIBLE_DEVICES"] = args.gpu if args.multi_gpu is None else '0,1,2,3'
'''CREATE DIR'''
experiment_dir = Path('./experiment/')
experiment_dir.mkdir(exist_ok=True)
file_dir = Path(
str(experiment_dir) + '/' +
str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')))
file_dir.mkdir(exist_ok=True)
checkpoints_dir = file_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = file_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger(args.model_name)
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler(
str(log_dir) + '/train_%s_partseg.txt' % args.model_name)
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
logger.info(
'---------------------------------------------------TRANING---------------------------------------------------'
)
logger.info('PARAMETER ...')
logger.info(args)
TRAIN_DATASET = PartNormalDataset(npoints=2048, split='trainval')
dataloader = torch.utils.data.DataLoader(TRAIN_DATASET,
batch_size=args.batchsize,
shuffle=True,
num_workers=int(args.workers))
TEST_DATASET = PartNormalDataset(npoints=2048, split='test')
testdataloader = torch.utils.data.DataLoader(TEST_DATASET,
batch_size=10,
shuffle=True,
num_workers=int(args.workers))
print("The number of training data is:", len(TRAIN_DATASET))
logger.info("The number of training data is:%d", len(TRAIN_DATASET))
print("The number of test data is:", len(TEST_DATASET))
logger.info("The number of test data is:%d", len(TEST_DATASET))
num_classes = 16
num_part = 50
blue = lambda x: '\033[94m' + x + '\033[0m'
model = PointNet2PartSeg_msg_one_hot(
num_part) if args.model_name == 'pointnet2' else PointNetDenseCls(
cat_num=num_classes, part_num=num_part)
if args.pretrain is not None:
model.load_state_dict(torch.load(args.pretrain))
print('load model %s' % args.pretrain)
logger.info('load model %s' % args.pretrain)
else:
print('Training from scratch')
logger.info('Training from scratch')
pretrain = args.pretrain
init_epoch = int(pretrain[-14:-11]) if args.pretrain is not None else 0
if args.optimizer == 'SGD':
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
elif args.optimizer == 'Adam':
optimizer = torch.optim.Adam(model.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=20,
gamma=0.5)
'''GPU selection and multi-GPU'''
if args.multi_gpu is not None:
device_ids = [int(x) for x in args.multi_gpu.split(',')]
torch.backends.cudnn.benchmark = True
model.cuda(device_ids[0])
model = torch.nn.DataParallel(model, device_ids=device_ids)
else:
model.cuda()
criterion = PointNetLoss()
LEARNING_RATE_CLIP = 1e-5
history = defaultdict(lambda: list())
best_acc = 0
best_class_avg_iou = 0
best_inctance_avg_iou = 0
for epoch in range(init_epoch, args.epoch):
scheduler.step()
lr = max(optimizer.param_groups[0]['lr'], LEARNING_RATE_CLIP)
print('Learning rate:%f' % lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
for i, data in tqdm(enumerate(dataloader, 0),
total=len(dataloader),
smoothing=0.9):
points, label, target, norm_plt = data
points, label, target = Variable(points.float()), Variable(
label.long()), Variable(target.long())
points = points.transpose(2, 1)
norm_plt = norm_plt.transpose(2, 1)
points, label, target, norm_plt = points.cuda(), label.squeeze(
).cuda(), target.cuda(), norm_plt.cuda()
optimizer.zero_grad()
model = model.train()
if args.model_name == 'pointnet':
labels_pred, seg_pred, trans_feat = model(
points, to_categorical(label, 16))
seg_pred = seg_pred.contiguous().view(-1, num_part)
target = target.view(-1, 1)[:, 0]
loss, seg_loss, label_loss = criterion(labels_pred, label,
seg_pred, target,
trans_feat)
else:
seg_pred = model(points, norm_plt, to_categorical(label, 16))
seg_pred = seg_pred.contiguous().view(-1, num_part)
target = target.view(-1, 1)[:, 0]
loss = F.nll_loss(seg_pred, target)
history['loss'].append(loss.cpu().data.numpy())
loss.backward()
optimizer.step()
forpointnet2 = args.model_name == 'pointnet2'
test_metrics, test_hist_acc, cat_mean_iou = test_partseg(
model, testdataloader, seg_label_to_cat, 50, forpointnet2)
print(
'Epoch %d %s accuracy: %f Class avg mIOU: %f Inctance avg mIOU: %f'
%
(epoch, blue('test'), test_metrics['accuracy'],
test_metrics['class_avg_iou'], test_metrics['inctance_avg_iou']))
logger.info(
'Epoch %d %s Accuracy: %f Class avg mIOU: %f Inctance avg mIOU: %f'
%
(epoch, blue('test'), test_metrics['accuracy'],
test_metrics['class_avg_iou'], test_metrics['inctance_avg_iou']))
if test_metrics['accuracy'] > best_acc:
best_acc = test_metrics['accuracy']
torch.save(
model.state_dict(), '%s/%s_%.3d_%.4f.pth' %
(checkpoints_dir, args.model_name, epoch, best_acc))
logger.info(cat_mean_iou)
logger.info('Save model..')
print('Save model..')
print(cat_mean_iou)
if test_metrics['class_avg_iou'] > best_class_avg_iou:
best_class_avg_iou = test_metrics['class_avg_iou']
if test_metrics['inctance_avg_iou'] > best_inctance_avg_iou:
best_inctance_avg_iou = test_metrics['inctance_avg_iou']
print('Best accuracy is: %.5f' % best_acc)
logger.info('Best accuracy is: %.5f' % best_acc)
print('Best class avg mIOU is: %.5f' % best_class_avg_iou)
logger.info('Best class avg mIOU is: %.5f' % best_class_avg_iou)
print('Best inctance avg mIOU is: %.5f' % best_inctance_avg_iou)
logger.info('Best inctance avg mIOU is: %.5f' % best_inctance_avg_iou)
def train(args):
experiment_dir = mkdir('./experiment/')
checkpoints_dir = mkdir('./experiment/partseg/%s/' % (args.model_name))
cache = _load(root)
norm = True if args.model_name == 'pointnet' else False
npoints = 2048
train_ds = PartNormalDataset(root,
cache,
npoints=npoints,
split='trainval',
data_augmentation=args.augment)
dataloader = DataLoader(train_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=int(args.workers))
test_ds = PartNormalDataset(root, cache, npoints=npoints, split='test')
testdataloader = DataLoader(test_ds,
batch_size=args.batch_size,
shuffle=False,
num_workers=int(args.workers))
num_classes = 16
num_part = 50
log.info(len_training=len(train_ds), len_testing=len(test_ds))
log.info(num_classes=num_classes, num_part=num_part)
if args.model_name == 'pointnet':
model = PointNetDenseCls(cat_num=num_classes, part_num=num_part)
else:
model = PointNet2PartSegMsg_one_hot(num_part)
torch.backends.cudnn.benchmark = True
model = torch.nn.DataParallel(model).cuda()
log.debug('Using gpu:', args.gpu)
if args.pretrain is not None and args.pretrain != 'None':
log.debug('Use pretrain model...')
model.load_state_dict(torch.load(args.pretrain))
init_epoch = int(args.pretrain[:-4].split('-')[-1])
log.debug('start epoch from', init_epoch)
else:
log.debug('Training from scratch')
init_epoch = 0
if args.optimizer == 'SGD':
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
elif args.optimizer == 'Adam':
optimizer = torch.optim.Adam(model.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=20,
gamma=0.5)
history = {'loss': []}
best_acc = 0
best_class_avg_iou = 0
best_inctance_avg_iou = 0
LEARNING_RATE_CLIP = 1e-5
# criterion = PointNetLoss()
def feature_transform_reguliarzer(trans):
d = trans.size()[1]
I = torch.eye(d)[None, :, :]
if trans.is_cuda:
I = I.cuda()
loss = torch.mean(
torch.norm(torch.bmm(trans,
trans.transpose(2, 1) - I), dim=(1, 2)))
return loss
def PointNet_Loss(labels_pred, label, seg_pred, seg, trans_feat):
mat_diff_loss_scale = 0.001
weight = 1
seg_loss = F.nll_loss(seg_pred, seg)
mat_diff_loss = feature_transform_reguliarzer(trans_feat)
label_loss = F.nll_loss(labels_pred, label)
loss = weight * seg_loss + (
1 - weight) * label_loss + mat_diff_loss * mat_diff_loss_scale
return loss, seg_loss, label_loss
for epoch in range(init_epoch, args.epoch):
scheduler.step()
lr = max(optimizer.param_groups[0]['lr'], LEARNING_RATE_CLIP)
log.info(job='partseg',
model=args.model_name,
gpu=args.gpu,
epoch='%d/%s' % (epoch, args.epoch),
lr=lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
for i, data in tqdm(enumerate(dataloader, 0),
total=len(dataloader),
smoothing=0.9):
points, label, target, norm_plt = data
points, label, target = points.float(), label.long(), target.long()
points = points.transpose(2, 1)
norm_plt = norm_plt.transpose(2, 1)
points, label, target, norm_plt = points.cuda(), label.squeeze(
).cuda(), target.cuda(), norm_plt.cuda()
optimizer.zero_grad()
model = model.train()
if args.model_name == 'pointnet':
labels_pred, seg_pred, trans_feat = model(
points, to_categorical(label, 16))
seg_pred = seg_pred.contiguous().view(-1, num_part)
target = target.view(-1, 1)[:, 0]
# loss, seg_loss, label_loss = criterion(labels_pred, label, seg_pred, target, trans_feat)
loss, seg_loss, label_loss = PointNet_Loss(
labels_pred, label, seg_pred, target, trans_feat)
else:
seg_pred = model(points, norm_plt, to_categorical(label, 16))
seg_pred = seg_pred.contiguous().view(-1, num_part)
target = target.view(-1, 1)[:, 0]
loss = F.nll_loss(seg_pred, target)
history['loss'].append(loss.cpu().data.numpy())
loss.backward()
optimizer.step()
log.debug('clear cuda cache')
torch.cuda.empty_cache()
test_metrics, test_hist_acc, cat_mean_iou = test_partseg(
model.eval(),
testdataloader,
label_id_to_name,
args.model_name,
num_part,
)
save_model = False
if test_metrics['accuracy'] > best_acc:
best_acc = test_metrics['accuracy']
if test_metrics['class_avg_iou'] > best_class_avg_iou:
best_class_avg_iou = test_metrics['class_avg_iou']
if test_metrics['inctance_avg_iou'] > best_inctance_avg_iou:
best_inctance_avg_iou = test_metrics['inctance_avg_iou']
save_model = True
if save_model:
fn_pth = 'partseg-%s-%.5f-%04d.pth' % (
args.model_name, best_inctance_avg_iou, epoch)
log.info('Save model...', fn=fn_pth)
torch.save(model.state_dict(),
os.path.join(checkpoints_dir, fn_pth))
log.info(cat_mean_iou)
else:
log.info('No need to save model')
log.warn('Curr',
accuracy=test_metrics['accuracy'],
class_avg_mIOU=test_metrics['class_avg_iou'],
inctance_avg_mIOU=test_metrics['inctance_avg_iou'])
log.warn('Best',
accuracy=best_acc,
class_avg_mIOU=best_class_avg_iou,
inctance_avg_mIOU=best_inctance_avg_iou)