def evaluate(args):
test_data, test_label = _load(load_train = False)
test_dataset = S3DISDataLoader(test_data,test_label)
testdataloader = DataLoader(test_dataset, batch_size=args.batch_size,shuffle=True, num_workers=args.workers)
log.debug('Building Model', args.model_name)
num_classes = 13
if args.model_name == 'pointnet':
model = PointNetSeg(num_classes, feature_transform=True, input_dims = 9)
else:
model = PointNet2SemSeg(num_classes, feature_dims = 6)
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)
test_metrics, cat_mean_iou = test_semseg(
model.eval(),
testdataloader,
label_id_to_name,
args.model_name,
num_classes,
)
mean_iou = np.mean(cat_mean_iou)
log.info(Test_accuracy=test_metrics['accuracy'], Test_meanIOU=mean_iou)
def main(args, stage=None, pretrain_model_path=None, e_init=0, e_final=2):
# Now load pickle labels mapping file
class_dict_fname = F_CLASS_DICT_PKL
print(class_dict_fname)
with open(class_dict_fname, "rb") as f:
class_dict = pickle.load(f)
f.close()
print("CLASS DICT: {}".format(class_dict))
# Use to get numeric classes --> semantic classes
seg_classes = class_dict
seg_label_to_cat = {}
for i, cat in enumerate(seg_classes.values()):
seg_label_to_cat[i] = cat
print('SEG LABEL', seg_label_to_cat)
# First load class weights file
with open(F_CLASS_WEIGHTS_PKL, "rb") as f:
class_weights = pickle.load(f)
f.close()
print('SEG CLASSES', seg_classes)
COUNTS = np.array(
[class_weights[key] for key in list(class_weights.keys())])
weight_normalizer = np.max(COUNTS)
weights = []
for count in COUNTS:
if count != 0:
weights.append(weight_normalizer / count)
else:
weights.append(0)
# Threshold weights
WEIGHTS_NP = np.array(weights)
WEIGHTS_NP[WEIGHTS_NP > THRESHOLD] = THRESHOLD
print("WEIGHTS ARE: {}".format(WEIGHTS_NP))
# Convert to pytorch tensor
weights = torch.from_numpy(WEIGHTS_NP.astype('float32'))
if USE_CLI:
gpu = args.gpu
multi_gpu = args.multi_gpu
batch_size = args.batch_size
model_name = args.model_name
optimizer = args.optimizer
learning_rate = args.learning_rate
pretrain = args.pretrain
multi_gpu = args.multi_gpu
batchsize = args.batchsize
decay_rate = args.decay_rate
epochs = args.epochs
else:
gpu = GPU
multi_gpu = MULTI_GPU
batch_size = BATCH_SIZE
model_name = MODEL_NAME
optimizer = OPTIMIZER
learning_rate = LEARNING_RATE
pretrain = PRETRAIN
multi_gpu = MULTI_GPU
batchsize = BATCH_SIZE
decay_rate = DECAY_RATE
epochs = EPOCHS
os.environ[
"CUDA_VISIBLE_DEVICES"] = gpu if multi_gpu is None else '0,1,2,3'
'''CREATE DIR'''
experiment_dir = Path('./experiment/{}'.format(EXPERIMENT_HEADER))
experiment_dir.mkdir(exist_ok=True)
file_dir = Path(
str(experiment_dir) + '/%sSemSeg-' % model_name +
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'''
if USE_CLI:
args = parse_args()
logger = logging.getLogger(model_name)
else:
logger = logging.getLogger(MODEL_NAME)
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
if USE_CLI:
file_handler = logging.FileHandler(
str(log_dir) + '/train_%s_semseg.txt' % args.model_name)
else:
file_handler = logging.FileHandler(
str(log_dir) + '/train_%s_semseg.txt' % MODEL_NAME)
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
logger.info(
'---------------------------------------------------TRANING---------------------------------------------------'
)
if USE_CLI:
logger.info('PARAMETER ...')
logger.info(args)
print('Load data...')
#train_data, train_label, test_data, test_label = recognize_all_data(test_area = 5)
# Now pickle our dataset
if USE_CLI:
f_in = args.data_path
else:
f_in = DATA_PATH
# Now pickle file
with open(f_in, "rb") as f:
DATA = pickle.load(f)
f.close()
# Now, we need to adjust labels according to our stage
labels_fpath = os.path.join("..", "data",
"stage_{}_labels.pkl".format(stage))
# Pickle this stage converter
with open(labels_fpath, "rb") as f:
labels = pickle.load(f)
print(type(labels), len(labels))
f.close()
# Now let's convert labels appropriately
keys = list(DATA.keys())
for key in keys:
DATA[key]['labels'] = labels[key]
random_seed = 42
indices = [i for i in range(len(list(DATA.keys())))]
np.random.seed(random_seed)
np.random.shuffle(indices)
TEST_SPLIT = 0.2
test_index = int(np.floor(TEST_SPLIT * len(list(DATA.keys()))))
print("val index is: {}".format(test_index))
train_indices, test_indices = indices[test_index:], indices[:test_index]
if USE_CLI:
print("LEN TRAIN: {}, LEN TEST: {}, EPOCHS: {}, OPTIMIZER: {}, DECAY_RATE: {}, LEARNING RATE: {}, \
DATA PATH: {}" .format(len(train_indices), len(test_indices), epochs, args.optimizer, args.decay_rate, \
args.learning_rate, args.data_path))
else:
print("LEN TRAIN: {}, LEN TEST: {}, EPOCHS: {}, OPTIMIZER: {}, DECAY_RATE: {}, LEARNING RATE: {}, \
DATA PATH: {}" .format(len(train_indices), len(test_indices), e_final-e_init, OPTIMIZER, DECAY_RATE, \
LEARNING_RATE, DATA_PATH))
# Creating PT data samplers and loaders:
train_sampler = SubsetRandomSampler(train_indices)
test_sampler = SubsetRandomSampler(test_indices)
print("INTERSECTION OF TRAIN/TEST (should be 0): {}".format(
len(set(train_indices).intersection(set(test_indices)))))
# Training dataset
dataset = A2D2DataLoader(DATA)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batchsize,
shuffle=False,
sampler=train_sampler,
collate_fn=collate_fn)
# Test dataset
test_dataset = A2D2DataLoader(DATA)
testdataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=batchsize,
shuffle=False,
sampler=test_sampler,
collate_fn=collate_fn)
num_classes = NUM_CLASSES
blue = lambda x: '\033[94m' + x + '\033[0m'
model = PointNet2SemSeg(
num_classes) if model_name == 'pointnet2' else PointNetSeg(
num_classes, feature_transform=True, semseg=True)
if pretrain_model_path is not None:
model.load_state_dict(torch.load(pretrain_model_path))
print('load model %s' % pretrain_model_path)
logger.info('load model %s' % pretrain_model_path)
else:
print('Training from scratch')
logger.info('Training from scratch')
#pretrain_var = pretrain
#init_epoch = int(pretrain_var[-14:-11]) if pretrain is not None else 0
init_epoch = e_init
if optimizer == 'SGD':
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
elif optimizer == 'Adam':
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=decay_rate)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=20,
gamma=0.5)
LEARNING_RATE_CLIP = 1e-5
'''GPU selection and multi-GPU'''
if multi_gpu is not None:
device_ids = [int(x) for x in 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()
history = defaultdict(lambda: list())
best_acc = 0
best_meaniou = 0
graph_losses = []
steps = []
step = 0
print("NUMBER OF EPOCHS IS: {}".format(e_final - e_init))
for epoch in range(e_init, e_final):
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
counter = 0
# Init confusion matrix
if USE_CONMAT:
conf_matrix = torch.zeros(NUM_CLASSES, NUM_CLASSES)
for points, targets in tqdm(dataloader):
#for points, target in tqdm(dataloader):
#points, target = data
points, targets = Variable(points.float()), Variable(
targets.long())
points = points.transpose(2, 1)
points, targets = points.cuda(), targets.cuda()
weights = weights.cuda()
optimizer.zero_grad() # REMOVE gradients
model = model.train()
if model_name == 'pointnet':
pred, trans_feat = model(points)
else:
pred = model(
points[:, :3, :], points[:, 3:, :]
) # Channels: xyz_norm (first 3) | rgb_norm (second three)
#pred = model(points)
if USE_CONMAT:
conf_matrix = confusion_matrix(pred, targets, conf_matrix)
pred = pred.contiguous().view(-1, num_classes)
targets = targets.view(-1, 1)[:, 0]
loss = F.nll_loss(pred, targets,
weight=weights) # Add class weights from dataset
if model_name == 'pointnet':
loss += feature_transform_reguliarzer(trans_feat) * 0.001
graph_losses.append(loss.cpu().data.numpy())
steps.append(step)
if counter % 100 == 0:
print("LOSS IS: {}".format(loss.cpu().data.numpy()))
#print((loss.cpu().data.numpy()))
history['loss'].append(loss.cpu().data.numpy())
loss.backward()
optimizer.step()
counter += 1
step += 1
#if counter > 3:
# break
if USE_CONMAT:
print("CONFUSION MATRIX: \n {}".format(conf_matrix))
pointnet2 = model_name == 'pointnet2'
test_metrics, test_hist_acc, cat_mean_iou = test_semseg(model.eval(), testdataloader, seg_label_to_cat,\
num_classes = num_classes,pointnet2=pointnet2)
mean_iou = np.mean(cat_mean_iou)
print('Epoch %d %s accuracy: %f meanIOU: %f' %
(epoch, blue('test'), test_metrics['accuracy'], mean_iou))
logger.info('Epoch %d %s accuracy: %f meanIOU: %f' %
(epoch, 'test', test_metrics['accuracy'], mean_iou))
if test_metrics['accuracy'] > best_acc:
best_acc = test_metrics['accuracy']
print("HERE")
save_path = '%s/%s_%.3d_%.4f_stage_%s.pth' % (
checkpoints_dir, model_name, epoch, best_acc, stage)
torch.save(model.state_dict(), save_path)
logger.info(cat_mean_iou)
logger.info('Save model..')
print('Save model..')
print(cat_mean_iou) #
if mean_iou > best_meaniou:
best_meaniou = mean_iou
print('Best accuracy is: %.5f' % best_acc)
logger.info('Best accuracy is: %.5f' % best_acc)
print('Best meanIOU is: %.5f' % best_meaniou)
logger.info('Best meanIOU is: %.5f' % best_meaniou)
if USE_CONMAT:
logger.info('Confusion matrix is: \n {}'.format(conf_matrix))
# Plot loss vs. steps
plt.plot(steps, graph_losses)
plt.xlabel("Batched Steps (Batch Size = {}".format(batch_size))
plt.ylabel("Multiclass NLL Loss")
plt.title("NLL Loss vs. Number of Batched Steps")
# Make directory for loss and other plots
graphs_dir = os.path.join(experiment_dir, "graphs")
os.makedirs(graphs_dir, exist_ok=True)
# Save and close figure
plt.savefig(os.path.join(graphs_dir, "losses.png"))
plt.clf()
return save_path
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) +'/%sSemSeg-'%args.model_name+ 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_semseg.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)
print('Load data...')
train_data, train_label, test_data, test_label = recognize_all_data(test_area = 5)
dataset = S3DISDataLoader(train_data,train_label)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batchsize,
shuffle=True, num_workers=int(args.workers))
test_dataset = S3DISDataLoader(test_data,test_label)
testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batchsize,
shuffle=True, num_workers=int(args.workers))
num_classes = 13
blue = lambda x: '\033[94m' + x + '\033[0m'
model = PointNet2SemSeg(num_classes) if args.model_name == 'pointnet2' else PointNetSeg(num_classes,feature_transform=True,semseg = True)
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)
LEARNING_RATE_CLIP = 1e-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()
history = defaultdict(lambda: list())
best_acc = 0
best_meaniou = 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, target = data
points, target = Variable(points.float()), Variable(target.long())
points = points.transpose(2, 1)
points, target = points.cuda(), target.cuda()
optimizer.zero_grad()
model = model.train()
if args.model_name == 'pointnet':
pred, trans_feat = model(points)
else:
pred = model(points[:,:3,:],points[:,3:,:])
pred = pred.contiguous().view(-1, num_classes)
target = target.view(-1, 1)[:, 0]
loss = F.nll_loss(pred, target)
if args.model_name == 'pointnet':
loss += feature_transform_reguliarzer(trans_feat) * 0.001
history['loss'].append(loss.cpu().data.numpy())
loss.backward()
optimizer.step()
pointnet2 = args.model_name == 'pointnet2'
test_metrics, test_hist_acc, cat_mean_iou = test_semseg(model.eval(), testdataloader, seg_label_to_cat,num_classes = num_classes,pointnet2=pointnet2)
mean_iou = np.mean(cat_mean_iou)
print('Epoch %d %s accuracy: %f meanIOU: %f' % (
epoch, blue('test'), test_metrics['accuracy'],mean_iou))
logger.info('Epoch %d %s accuracy: %f meanIOU: %f' % (
epoch, 'test', test_metrics['accuracy'],mean_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 mean_iou > best_meaniou:
best_meaniou = mean_iou
print('Best accuracy is: %.5f'%best_acc)
logger.info('Best accuracy is: %.5f'%best_acc)
print('Best meanIOU is: %.5f'%best_meaniou)
logger.info('Best meanIOU is: %.5f'%best_meaniou)
def train(args):
experiment_dir = mkdir('./experiment/')
checkpoints_dir = mkdir('./experiment/semseg/%s/'%(args.model_name))
train_data, train_label, test_data, test_label = _load()
dataset = S3DISDataLoader(train_data, train_label, data_augmentation = args.augment)
dataloader = DataLoader(dataset, batch_size=args.batch_size,shuffle=True, num_workers=args.workers)
test_dataset = S3DISDataLoader(test_data, test_label)
testdataloader = DataLoader(test_dataset, batch_size=args.batch_size,shuffle=True, num_workers=args.workers)
num_classes = 13
if args.model_name == 'pointnet':
model = PointNetSeg(num_classes, feature_transform=True, input_dims = 9)
else:
model = PointNet2SemSeg(num_classes, feature_dims = 6)
torch.backends.cudnn.benchmark = True
model = torch.nn.DataParallel(model).cuda()
log.debug('Using gpu:',args.gpu)
if args.pretrain is not 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)
LEARNING_RATE_CLIP = 1e-5
history = {'loss':[]}
best_acc = 0
best_meaniou = 0
for epoch in range(init_epoch,args.epoch):
scheduler.step()
lr = max(optimizer.param_groups[0]['lr'],LEARNING_RATE_CLIP)
log.info(job='semseg',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 points, target in tqdm(dataloader, total=len(dataloader), smoothing=0.9, dynamic_ncols=True):
points, target = points.float(), target.long()
points = points.transpose(2, 1)
points, target = points.cuda(), target.cuda()
optimizer.zero_grad()
model = model.train()
if args.model_name == 'pointnet':
pred, trans_feat = model(points)
else:
pred = model(points)
pred = pred.contiguous().view(-1, num_classes)
target = target.view(-1, 1)[:, 0]
loss = F.nll_loss(pred, target)
if args.model_name == 'pointnet':
loss += feature_transform_reguliarzer(trans_feat) * 0.001
history['loss'].append(loss.cpu().data.numpy())
loss.backward()
optimizer.step()
log.debug('clear cuda cache')
torch.cuda.empty_cache()
test_metrics, cat_mean_iou = test_semseg(
model.eval(),
testdataloader,
label_id_to_name,
args.model_name,
num_classes,
)
mean_iou = np.mean(cat_mean_iou)
save_model = False
if test_metrics['accuracy'] > best_acc:
best_acc = test_metrics['accuracy']
if mean_iou > best_meaniou:
best_meaniou = mean_iou
save_model = True
if save_model:
fn_pth = 'semseg-%s-%.5f-%04d.pth' % (args.model_name, best_meaniou, epoch)
log.info('Save model...',fn = fn_pth)
torch.save(model.state_dict(), os.path.join(checkpoints_dir, fn_pth))
log.warn(cat_mean_iou)
else:
log.info('No need to save model')
log.warn(cat_mean_iou)
log.warn('Curr',accuracy=test_metrics['accuracy'], meanIOU=mean_iou)
log.warn('Best',accuracy=best_acc, meanIOU=best_meaniou)