From 88790ddf234b5fb0242e55d78c281524460a999f Mon Sep 17 00:00:00 2001 From: Benny <775410794@qq.com> Date: Sat, 20 Mar 2021 23:38:12 +0800 Subject: [PATCH] re-organize codes --- data_utils/S3DISDataLoader.py | 6 +- test_partseg.py | 50 +++++++++------ test_semseg.py | 35 ++++++---- train_partseg.py | 116 +++++++++++++++++++--------------- train_semseg.py | 66 +++++++++++-------- 5 files changed, 161 insertions(+), 112 deletions(-) diff --git a/data_utils/S3DISDataLoader.py b/data_utils/S3DISDataLoader.py index 5f6602e5b..d4aae42e6 100644 --- a/data_utils/S3DISDataLoader.py +++ b/data_utils/S3DISDataLoader.py @@ -1,5 +1,7 @@ import os import numpy as np + +from tqdm import tqdm from torch.utils.data import Dataset @@ -15,11 +17,13 @@ def __init__(self, split='train', data_root='trainval_fullarea', num_point=4096, rooms_split = [room for room in rooms if not 'Area_{}'.format(test_area) in room] else: rooms_split = [room for room in rooms if 'Area_{}'.format(test_area) in room] + self.room_points, self.room_labels = [], [] self.room_coord_min, self.room_coord_max = [], [] num_point_all = [] labelweights = np.zeros(13) - for room_name in rooms_split: + + for room_name in tqdm(rooms_split, total=len(rooms_split)): room_path = os.path.join(data_root, room_name) room_data = np.load(room_path) # xyzrgbl, N*7 points, labels = room_data[:, 0:6], room_data[:, 6] # xyzrgb, N*6; l, N diff --git a/test_partseg.py b/test_partseg.py index 122791c21..687e1deab 100644 --- a/test_partseg.py +++ b/test_partseg.py @@ -16,12 +16,17 @@ ROOT_DIR = BASE_DIR sys.path.append(os.path.join(ROOT_DIR, 'models')) -seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]} -seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} +seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], + 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37], + 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49], + 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]} + +seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} for cat in seg_classes.keys(): for label in seg_classes[cat]: seg_label_to_cat[label] = cat + def to_categorical(y, num_classes): """ 1-hot encodes a tensor """ new_y = torch.eye(num_classes)[y.cpu().data.numpy(),] @@ -33,14 +38,15 @@ def to_categorical(y, num_classes): def parse_args(): '''PARAMETERS''' parser = argparse.ArgumentParser('PointNet') - parser.add_argument('--batch_size', type=int, default=24, help='batch size in testing [default: 24]') - parser.add_argument('--gpu', type=str, default='0', help='specify gpu device [default: 0]') - parser.add_argument('--num_point', type=int, default=2048, help='Point Number [default: 2048]') - parser.add_argument('--log_dir', type=str, default='pointnet2_part_seg_ssg', help='Experiment root') - parser.add_argument('--normal', action='store_true', default=False, help='Whether to use normal information [default: False]') - parser.add_argument('--num_votes', type=int, default=3, help='Aggregate segmentation scores with voting [default: 3]') + parser.add_argument('--batch_size', type=int, default=24, help='batch size in testing') + parser.add_argument('--gpu', type=str, default='0', help='specify gpu device') + parser.add_argument('--num_point', type=int, default=2048, help='point Number') + parser.add_argument('--log_dir', type=str, required=True, help='experiment root') + parser.add_argument('--normal', action='store_true', default=False, help='use normals') + parser.add_argument('--num_votes', type=int, default=3, help='aggregate segmentation scores with voting') return parser.parse_args() + def main(args): def log_string(str): logger.info(str) @@ -64,20 +70,19 @@ def log_string(str): root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/' - TEST_DATASET = PartNormalDataset(root = root, npoints=args.num_point, split='test', normal_channel=args.normal) - testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size,shuffle=False, num_workers=4) - log_string("The number of test data is: %d" % len(TEST_DATASET)) + TEST_DATASET = PartNormalDataset(root=root, npoints=args.num_point, split='test', normal_channel=args.normal) + testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) + log_string("The number of test data is: %d" % len(TEST_DATASET)) num_classes = 16 num_part = 50 '''MODEL LOADING''' - model_name = os.listdir(experiment_dir+'/logs')[0].split('.')[0] + model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0] MODEL = importlib.import_module(model_name) classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda() checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth') classifier.load_state_dict(checkpoint['model_state_dict']) - with torch.no_grad(): test_metrics = {} total_correct = 0 @@ -86,29 +91,35 @@ def log_string(str): total_correct_class = [0 for _ in range(num_part)] shape_ious = {cat: [] for cat in seg_classes.keys()} seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} + for cat in seg_classes.keys(): for label in seg_classes[cat]: seg_label_to_cat[label] = cat - for batch_id, (points, label, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9): + classifier = classifier.eval() + for batch_id, (points, label, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), + smoothing=0.9): batchsize, num_point, _ = points.size() cur_batch_size, NUM_POINT, _ = points.size() points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda() points = points.transpose(2, 1) - classifier = classifier.eval() vote_pool = torch.zeros(target.size()[0], target.size()[1], num_part).cuda() + for _ in range(args.num_votes): seg_pred, _ = classifier(points, to_categorical(label, num_classes)) vote_pool += seg_pred + seg_pred = vote_pool / args.num_votes cur_pred_val = seg_pred.cpu().data.numpy() cur_pred_val_logits = cur_pred_val cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32) target = target.cpu().data.numpy() + for i in range(cur_batch_size): cat = seg_label_to_cat[target[i, 0]] logits = cur_pred_val_logits[i, :, :] cur_pred_val[i, :] = np.argmax(logits[:, seg_classes[cat]], 1) + seg_classes[cat][0] + correct = np.sum(cur_pred_val == target) total_correct += correct total_seen += (cur_batch_size * NUM_POINT) @@ -145,13 +156,12 @@ def log_string(str): test_metrics['class_avg_iou'] = mean_shape_ious test_metrics['inctance_avg_iou'] = np.mean(all_shape_ious) + log_string('Accuracy is: %.5f' % test_metrics['accuracy']) + log_string('Class avg accuracy is: %.5f' % test_metrics['class_avg_accuracy']) + log_string('Class avg mIOU is: %.5f' % test_metrics['class_avg_iou']) + log_string('Inctance avg mIOU is: %.5f' % test_metrics['inctance_avg_iou']) - log_string('Accuracy is: %.5f'%test_metrics['accuracy']) - log_string('Class avg accuracy is: %.5f'%test_metrics['class_avg_accuracy']) - log_string('Class avg mIOU is: %.5f'%test_metrics['class_avg_iou']) - log_string('Inctance avg mIOU is: %.5f'%test_metrics['inctance_avg_iou']) if __name__ == '__main__': args = parse_args() main(args) - diff --git a/test_semseg.py b/test_semseg.py index 4541b0a95..3df67ac94 100644 --- a/test_semseg.py +++ b/test_semseg.py @@ -19,13 +19,15 @@ ROOT_DIR = BASE_DIR sys.path.append(os.path.join(ROOT_DIR, 'models')) -classes = ['ceiling','floor','wall','beam','column','window','door','table','chair','sofa','bookcase','board','clutter'] -class2label = {cls: i for i,cls in enumerate(classes)} +classes = ['ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door', 'table', 'chair', 'sofa', 'bookcase', + 'board', 'clutter'] +class2label = {cls: i for i, cls in enumerate(classes)} seg_classes = class2label seg_label_to_cat = {} -for i,cat in enumerate(seg_classes.keys()): +for i, cat in enumerate(seg_classes.keys()): seg_label_to_cat[i] = cat + def parse_args(): '''PARAMETERS''' parser = argparse.ArgumentParser('Model') @@ -38,15 +40,17 @@ def parse_args(): parser.add_argument('--num_votes', type=int, default=5, help='Aggregate segmentation scores with voting [default: 5]') return parser.parse_args() + def add_vote(vote_label_pool, point_idx, pred_label, weight): B = pred_label.shape[0] N = pred_label.shape[1] for b in range(B): for n in range(N): - if weight[b,n]: + if weight[b, n]: vote_label_pool[int(point_idx[b, n]), int(pred_label[b, n])] += 1 return vote_label_pool + def main(args): def log_string(str): logger.info(str) @@ -75,17 +79,18 @@ def log_string(str): BATCH_SIZE = args.batch_size NUM_POINT = args.num_point - root = 'data/stanford_indoor3d/' + root = 'data/s3dis/stanford_indoor3d/' TEST_DATASET_WHOLE_SCENE = ScannetDatasetWholeScene(root, split='test', test_area=args.test_area, block_points=NUM_POINT) - log_string("The number of test data is: %d" % len(TEST_DATASET_WHOLE_SCENE)) + log_string("The number of test data is: %d" % len(TEST_DATASET_WHOLE_SCENE)) '''MODEL LOADING''' - model_name = os.listdir(experiment_dir+'/logs')[0].split('.')[0] + model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0] MODEL = importlib.import_module(model_name) classifier = MODEL.get_model(NUM_CLASSES).cuda() checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth') classifier.load_state_dict(checkpoint['model_state_dict']) + classifier = classifier.eval() with torch.no_grad(): scene_id = TEST_DATASET_WHOLE_SCENE.file_list @@ -99,7 +104,7 @@ def log_string(str): log_string('---- EVALUATION WHOLE SCENE----') for batch_idx in range(num_batches): - print("visualize [%d/%d] %s ..." % (batch_idx+1, num_batches, scene_id[batch_idx])) + print("visualize [%d/%d] %s ..." % (batch_idx + 1, num_batches, scene_id[batch_idx])) total_seen_class_tmp = [0 for _ in range(NUM_CLASSES)] total_correct_class_tmp = [0 for _ in range(NUM_CLASSES)] total_iou_deno_class_tmp = [0 for _ in range(NUM_CLASSES)] @@ -119,6 +124,7 @@ def log_string(str): batch_label = np.zeros((BATCH_SIZE, NUM_POINT)) batch_point_index = np.zeros((BATCH_SIZE, NUM_POINT)) batch_smpw = np.zeros((BATCH_SIZE, NUM_POINT)) + for sbatch in range(s_batch_num): start_idx = sbatch * BATCH_SIZE end_idx = min((sbatch + 1) * BATCH_SIZE, num_blocks) @@ -130,7 +136,7 @@ def log_string(str): batch_data[:, :, 3:6] /= 1.0 torch_data = torch.Tensor(batch_data) - torch_data= torch_data.float().cuda() + torch_data = torch_data.float().cuda() torch_data = torch_data.transpose(2, 1) seg_pred, _ = classifier(torch_data) batch_pred_label = seg_pred.contiguous().cpu().data.max(2)[1].numpy() @@ -166,12 +172,12 @@ def log_string(str): color_gt = g_label2color[whole_scene_label[i]] if args.visual: fout.write('v %f %f %f %d %d %d\n' % ( - whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color[0], color[1], - color[2])) + whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color[0], color[1], + color[2])) fout_gt.write( 'v %f %f %f %d %d %d\n' % ( - whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color_gt[0], - color_gt[1], color_gt[2])) + whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color_gt[0], + color_gt[1], color_gt[2])) if args.visual: fout.close() fout_gt.close() @@ -187,10 +193,11 @@ def log_string(str): log_string('eval whole scene point avg class acc: %f' % ( np.mean(np.array(total_correct_class) / (np.array(total_seen_class, dtype=np.float) + 1e-6)))) log_string('eval whole scene point accuracy: %f' % ( - np.sum(total_correct_class) / float(np.sum(total_seen_class) + 1e-6))) + np.sum(total_correct_class) / float(np.sum(total_seen_class) + 1e-6))) print("Done!") + if __name__ == '__main__': args = parse_args() main(args) diff --git a/train_partseg.py b/train_partseg.py index 38bcbe63b..6924fa857 100644 --- a/train_partseg.py +++ b/train_partseg.py @@ -20,12 +20,16 @@ ROOT_DIR = BASE_DIR sys.path.append(os.path.join(ROOT_DIR, 'models')) -seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]} -seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} +seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], + 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37], + 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49], + 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]} +seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} for cat in seg_classes.keys(): for label in seg_classes[cat]: seg_label_to_cat[label] = cat + def to_categorical(y, num_classes): """ 1-hot encodes a tensor """ new_y = torch.eye(num_classes)[y.cpu().data.numpy(),] @@ -36,21 +40,22 @@ def to_categorical(y, num_classes): def parse_args(): parser = argparse.ArgumentParser('Model') - parser.add_argument('--model', type=str, default='pointnet2_part_seg_msg', help='model name [default: pointnet2_part_seg_msg]') - parser.add_argument('--batch_size', type=int, default=16, help='Batch Size during training [default: 16]') - parser.add_argument('--epoch', default=251, type=int, help='Epoch to run [default: 251]') - parser.add_argument('--learning_rate', default=0.001, type=float, help='Initial learning rate [default: 0.001]') - parser.add_argument('--gpu', type=str, default='0', help='GPU to use [default: GPU 0]') - parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD [default: Adam]') - parser.add_argument('--log_dir', type=str, default=None, help='Log path [default: None]') - parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay [default: 1e-4]') - parser.add_argument('--npoint', type=int, default=2048, help='Point Number [default: 2048]') - parser.add_argument('--normal', action='store_true', default=False, help='Whether to use normal information [default: False]') - parser.add_argument('--step_size', type=int, default=20, help='Decay step for lr decay [default: every 20 epochs]') - parser.add_argument('--lr_decay', type=float, default=0.5, help='Decay rate for lr decay [default: 0.5]') + parser.add_argument('--model', type=str, default='pointnet2_part_seg_msg', help='model name') + parser.add_argument('--batch_size', type=int, default=16, help='batch Size during training') + parser.add_argument('--epoch', default=251, type=int, help='epoch to run') + parser.add_argument('--learning_rate', default=0.001, type=float, help='initial learning rate') + parser.add_argument('--gpu', type=str, default='0', help='specify GPU devices') + parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD') + parser.add_argument('--log_dir', type=str, default=None, help='log path') + parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay') + parser.add_argument('--npoint', type=int, default=2048, help='point Number') + parser.add_argument('--normal', action='store_true', default=False, help='use normals') + parser.add_argument('--step_size', type=int, default=20, help='decay step for lr decay') + parser.add_argument('--lr_decay', type=float, default=0.5, help='decay rate for lr decay') return parser.parse_args() + def main(args): def log_string(str): logger.info(str) @@ -61,18 +66,18 @@ def log_string(str): '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) - experiment_dir = Path('./log/') - experiment_dir.mkdir(exist_ok=True) - experiment_dir = experiment_dir.joinpath('part_seg') - experiment_dir.mkdir(exist_ok=True) + exp_dir = Path('./log/') + exp_dir.mkdir(exist_ok=True) + exp_dir = exp_dir.joinpath('part_seg') + exp_dir.mkdir(exist_ok=True) if args.log_dir is None: - experiment_dir = experiment_dir.joinpath(timestr) + exp_dir = exp_dir.joinpath(timestr) else: - experiment_dir = experiment_dir.joinpath(args.log_dir) - experiment_dir.mkdir(exist_ok=True) - checkpoints_dir = experiment_dir.joinpath('checkpoints/') + exp_dir = exp_dir.joinpath(args.log_dir) + exp_dir.mkdir(exist_ok=True) + checkpoints_dir = exp_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) - log_dir = experiment_dir.joinpath('logs/') + log_dir = exp_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' @@ -89,23 +94,25 @@ def log_string(str): root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/' - TRAIN_DATASET = PartNormalDataset(root = root, npoints=args.npoint, split='trainval', normal_channel=args.normal) - trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size,shuffle=True, num_workers=4) - TEST_DATASET = PartNormalDataset(root = root, npoints=args.npoint, split='test', normal_channel=args.normal) - testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size,shuffle=False, num_workers=4) + TRAIN_DATASET = PartNormalDataset(root=root, npoints=args.npoint, split='trainval', normal_channel=args.normal) + trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, + num_workers=4) + TEST_DATASET = PartNormalDataset(root=root, npoints=args.npoint, split='test', normal_channel=args.normal) + testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) log_string("The number of training data is: %d" % len(TRAIN_DATASET)) - log_string("The number of test data is: %d" % len(TEST_DATASET)) + log_string("The number of test data is: %d" % len(TEST_DATASET)) + num_classes = 16 num_part = 50 + '''MODEL LOADING''' MODEL = importlib.import_module(args.model) - shutil.copy('models/%s.py' % args.model, str(experiment_dir)) - shutil.copy('models/pointnet2_utils.py', str(experiment_dir)) + shutil.copy('models/%s.py' % args.model, str(exp_dir)) + shutil.copy('models/pointnet2_utils.py', str(exp_dir)) classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() - def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv2d') != -1: @@ -116,7 +123,7 @@ def weights_init(m): torch.nn.init.constant_(m.bias.data, 0.0) try: - checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth') + checkpoint = torch.load(str(exp_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') @@ -150,40 +157,44 @@ def bn_momentum_adjust(m, momentum): best_class_avg_iou = 0 best_inctance_avg_iou = 0 - for epoch in range(start_epoch,args.epoch): + for epoch in range(start_epoch, args.epoch): + mean_correct = [] + log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) '''Adjust learning rate and BN momentum''' lr = max(args.learning_rate * (args.lr_decay ** (epoch // args.step_size)), LEARNING_RATE_CLIP) log_string('Learning rate:%f' % lr) for param_group in optimizer.param_groups: param_group['lr'] = lr - mean_correct = [] momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY ** (epoch // MOMENTUM_DECCAY_STEP)) if momentum < 0.01: momentum = 0.01 print('BN momentum updated to: %f' % momentum) - classifier = classifier.apply(lambda x: bn_momentum_adjust(x,momentum)) + classifier = classifier.apply(lambda x: bn_momentum_adjust(x, momentum)) + classifier = classifier.train() '''learning one epoch''' - for i, data in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9): - points, label, target = data + for i, (points, label, target) in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9): + optimizer.zero_grad() + points = points.data.numpy() - points[:,:, 0:3] = provider.random_scale_point_cloud(points[:,:, 0:3]) - points[:,:, 0:3] = provider.shift_point_cloud(points[:,:, 0:3]) + points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) + points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) points = torch.Tensor(points) - points, label, target = points.float().cuda(),label.long().cuda(), target.long().cuda() + points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda() points = points.transpose(2, 1) - optimizer.zero_grad() - classifier = classifier.train() + seg_pred, trans_feat = classifier(points, to_categorical(label, num_classes)) seg_pred = seg_pred.contiguous().view(-1, num_part) target = target.view(-1, 1)[:, 0] pred_choice = seg_pred.data.max(1)[1] + correct = pred_choice.eq(target.data).cpu().sum() mean_correct.append(correct.item() / (args.batch_size * args.npoint)) loss = criterion(seg_pred, target, trans_feat) loss.backward() optimizer.step() + train_instance_acc = np.mean(mean_correct) log_string('Train accuracy is: %.5f' % train_instance_acc) @@ -195,24 +206,28 @@ def bn_momentum_adjust(m, momentum): total_correct_class = [0 for _ in range(num_part)] shape_ious = {cat: [] for cat in seg_classes.keys()} seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} + for cat in seg_classes.keys(): for label in seg_classes[cat]: seg_label_to_cat[label] = cat + classifier = classifier.eval() + for batch_id, (points, label, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9): cur_batch_size, NUM_POINT, _ = points.size() points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda() points = points.transpose(2, 1) - classifier = classifier.eval() seg_pred, _ = classifier(points, to_categorical(label, num_classes)) cur_pred_val = seg_pred.cpu().data.numpy() cur_pred_val_logits = cur_pred_val cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32) target = target.cpu().data.numpy() + for i in range(cur_batch_size): cat = seg_label_to_cat[target[i, 0]] logits = cur_pred_val_logits[i, :, :] cur_pred_val[i, :] = np.argmax(logits[:, seg_classes[cat]], 1) + seg_classes[cat][0] + correct = np.sum(cur_pred_val == target) total_correct += correct total_seen += (cur_batch_size * NUM_POINT) @@ -249,13 +264,12 @@ def bn_momentum_adjust(m, momentum): test_metrics['class_avg_iou'] = mean_shape_ious test_metrics['inctance_avg_iou'] = np.mean(all_shape_ious) - log_string('Epoch %d test Accuracy: %f Class avg mIOU: %f Inctance avg mIOU: %f' % ( - epoch+1, test_metrics['accuracy'],test_metrics['class_avg_iou'],test_metrics['inctance_avg_iou'])) + epoch + 1, test_metrics['accuracy'], test_metrics['class_avg_iou'], test_metrics['inctance_avg_iou'])) if (test_metrics['inctance_avg_iou'] >= best_inctance_avg_iou): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' - log_string('Saving at %s'% savepath) + log_string('Saving at %s' % savepath) state = { 'epoch': epoch, 'train_acc': train_instance_acc, @@ -274,12 +288,12 @@ def bn_momentum_adjust(m, momentum): 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'] - log_string('Best accuracy is: %.5f'%best_acc) - log_string('Best class avg mIOU is: %.5f'%best_class_avg_iou) - log_string('Best inctance avg mIOU is: %.5f'%best_inctance_avg_iou) - global_epoch+=1 + log_string('Best accuracy is: %.5f' % best_acc) + log_string('Best class avg mIOU is: %.5f' % best_class_avg_iou) + log_string('Best inctance avg mIOU is: %.5f' % best_inctance_avg_iou) + global_epoch += 1 + if __name__ == '__main__': args = parse_args() main(args) - diff --git a/train_semseg.py b/train_semseg.py index 1e0f329c8..0f8c10f6e 100644 --- a/train_semseg.py +++ b/train_semseg.py @@ -21,12 +21,12 @@ ROOT_DIR = BASE_DIR sys.path.append(os.path.join(ROOT_DIR, 'models')) - -classes = ['ceiling','floor','wall','beam','column','window','door','table','chair','sofa','bookcase','board','clutter'] -class2label = {cls: i for i,cls in enumerate(classes)} +classes = ['ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door', 'table', 'chair', 'sofa', 'bookcase', + 'board', 'clutter'] +class2label = {cls: i for i, cls in enumerate(classes)} seg_classes = class2label seg_label_to_cat = {} -for i,cat in enumerate(seg_classes.keys()): +for i, cat in enumerate(seg_classes.keys()): seg_label_to_cat[i] = cat @@ -34,19 +34,20 @@ def parse_args(): parser = argparse.ArgumentParser('Model') parser.add_argument('--model', type=str, default='pointnet_sem_seg', help='model name [default: pointnet_sem_seg]') parser.add_argument('--batch_size', type=int, default=16, help='Batch Size during training [default: 16]') - parser.add_argument('--epoch', default=128, type=int, help='Epoch to run [default: 128]') + parser.add_argument('--epoch', default=128, type=int, help='Epoch to run [default: 128]') parser.add_argument('--learning_rate', default=0.001, type=float, help='Initial learning rate [default: 0.001]') parser.add_argument('--gpu', type=str, default='0', help='GPU to use [default: GPU 0]') parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD [default: Adam]') parser.add_argument('--log_dir', type=str, default=None, help='Log path [default: None]') parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay [default: 1e-4]') - parser.add_argument('--npoint', type=int, default=4096, help='Point Number [default: 4096]') - parser.add_argument('--step_size', type=int, default=10, help='Decay step for lr decay [default: every 10 epochs]') - parser.add_argument('--lr_decay', type=float, default=0.7, help='Decay rate for lr decay [default: 0.7]') + parser.add_argument('--npoint', type=int, default=4096, help='Point Number [default: 4096]') + parser.add_argument('--step_size', type=int, default=10, help='Decay step for lr decay [default: every 10 epochs]') + parser.add_argument('--lr_decay', type=float, default=0.7, help='Decay rate for lr decay [default: 0.7]') parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') return parser.parse_args() + def main(args): def log_string(str): logger.info(str) @@ -83,7 +84,7 @@ def log_string(str): log_string('PARAMETER ...') log_string(args) - root = 'data/stanford_indoor3d/' + root = 'data/s3dis/stanford_indoor3d/' NUM_CLASSES = 13 NUM_POINT = args.npoint BATCH_SIZE = args.batch_size @@ -92,8 +93,12 @@ def log_string(str): TRAIN_DATASET = S3DISDataset(split='train', data_root=root, num_point=NUM_POINT, test_area=args.test_area, block_size=1.0, sample_rate=1.0, transform=None) print("start loading test data ...") TEST_DATASET = S3DISDataset(split='test', data_root=root, num_point=NUM_POINT, test_area=args.test_area, block_size=1.0, sample_rate=1.0, transform=None) - trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, worker_init_fn = lambda x: np.random.seed(x+int(time.time()))) - testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True, drop_last=True) + + trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=BATCH_SIZE, shuffle=True, num_workers=10, + pin_memory=True, drop_last=True, + worker_init_fn=lambda x: np.random.seed(x + int(time.time()))) + testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=BATCH_SIZE, shuffle=False, num_workers=10, + pin_memory=True, drop_last=True) weights = torch.Tensor(TRAIN_DATASET.labelweights).cuda() log_string("The number of training data is: %d" % len(TRAIN_DATASET)) @@ -149,7 +154,7 @@ def bn_momentum_adjust(m, momentum): global_epoch = 0 best_iou = 0 - for epoch in range(start_epoch,args.epoch): + for epoch in range(start_epoch, args.epoch): '''Train on chopped scenes''' log_string('**** Epoch %d (%d/%s) ****' % (global_epoch + 1, epoch + 1, args.epoch)) lr = max(args.learning_rate * (args.lr_decay ** (epoch // args.step_size)), LEARNING_RATE_CLIP) @@ -160,27 +165,31 @@ def bn_momentum_adjust(m, momentum): if momentum < 0.01: momentum = 0.01 print('BN momentum updated to: %f' % momentum) - classifier = classifier.apply(lambda x: bn_momentum_adjust(x,momentum)) + classifier = classifier.apply(lambda x: bn_momentum_adjust(x, momentum)) num_batches = len(trainDataLoader) total_correct = 0 total_seen = 0 loss_sum = 0 - for i, data in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9): - points, target = data + classifier = classifier.train() + + for i, (points, target) in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9): + optimizer.zero_grad() + points = points.data.numpy() - points[:,:, :3] = provider.rotate_point_cloud_z(points[:,:, :3]) + points[:, :, :3] = provider.rotate_point_cloud_z(points[:, :, :3]) points = torch.Tensor(points) - points, target = points.float().cuda(),target.long().cuda() + points, target = points.float().cuda(), target.long().cuda() points = points.transpose(2, 1) - optimizer.zero_grad() - classifier = classifier.train() + seg_pred, trans_feat = classifier(points) seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES) + batch_label = target.view(-1, 1)[:, 0].cpu().data.numpy() target = target.view(-1, 1)[:, 0] loss = criterion(seg_pred, target, trans_feat, weights) loss.backward() optimizer.step() + pred_choice = seg_pred.cpu().data.max(1)[1].numpy() correct = np.sum(pred_choice == batch_label) total_correct += correct @@ -211,17 +220,19 @@ def bn_momentum_adjust(m, momentum): total_seen_class = [0 for _ in range(NUM_CLASSES)] total_correct_class = [0 for _ in range(NUM_CLASSES)] total_iou_deno_class = [0 for _ in range(NUM_CLASSES)] + classifier = classifier.eval() + log_string('---- EPOCH %03d EVALUATION ----' % (global_epoch + 1)) - for i, data in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9): - points, target = data + for i, (points, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9): points = points.data.numpy() points = torch.Tensor(points) points, target = points.float().cuda(), target.long().cuda() points = points.transpose(2, 1) - classifier = classifier.eval() + seg_pred, trans_feat = classifier(points) pred_val = seg_pred.contiguous().cpu().data.numpy() seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES) + batch_label = target.cpu().data.numpy() target = target.view(-1, 1)[:, 0] loss = criterion(seg_pred, target, trans_feat, weights) @@ -232,10 +243,12 @@ def bn_momentum_adjust(m, momentum): total_seen += (BATCH_SIZE * NUM_POINT) tmp, _ = np.histogram(batch_label, range(NUM_CLASSES + 1)) labelweights += tmp + for l in range(NUM_CLASSES): - total_seen_class[l] += np.sum((batch_label == l) ) - total_correct_class[l] += np.sum((pred_val == l) & (batch_label == l) ) - total_iou_deno_class[l] += np.sum(((pred_val == l) | (batch_label == l)) ) + total_seen_class[l] += np.sum((batch_label == l)) + total_correct_class[l] += np.sum((pred_val == l) & (batch_label == l)) + total_iou_deno_class[l] += np.sum(((pred_val == l) | (batch_label == l))) + labelweights = labelweights.astype(np.float32) / np.sum(labelweights.astype(np.float32)) mIoU = np.mean(np.array(total_correct_class) / (np.array(total_iou_deno_class, dtype=np.float) + 1e-6)) log_string('eval mean loss: %f' % (loss_sum / float(num_batches))) @@ -243,6 +256,7 @@ def bn_momentum_adjust(m, momentum): log_string('eval point accuracy: %f' % (total_correct / float(total_seen))) log_string('eval point avg class acc: %f' % ( np.mean(np.array(total_correct_class) / (np.array(total_seen_class, dtype=np.float) + 1e-6)))) + iou_per_class_str = '------- IoU --------\n' for l in range(NUM_CLASSES): iou_per_class_str += 'class %s weight: %.3f, IoU: %.3f \n' % ( @@ -252,6 +266,7 @@ def bn_momentum_adjust(m, momentum): log_string(iou_per_class_str) log_string('Eval mean loss: %f' % (loss_sum / num_batches)) log_string('Eval accuracy: %f' % (total_correct / float(total_seen))) + if mIoU >= best_iou: best_iou = mIoU logger.info('Save model...') @@ -272,4 +287,3 @@ def bn_momentum_adjust(m, momentum): if __name__ == '__main__': args = parse_args() main(args) -