def main(args, GAMMA, pretrain_model_path=None): # 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 = {} print(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.format(GAMMA))) 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() 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), EPOCHS, 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 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(epochs)) for epoch in range(epochs): 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 = FocalLoss(gamma=GAMMA)(pred, targets) #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.pth' % (checkpoints_dir, model_name, epoch, best_acc) 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()
args = argparse.Namespace() cuda_available = torch.cuda.is_available() if cuda_available: torch.cuda.set_device(0) device = torch.device("cuda" if cuda_available else "cpu") args.device = device plt.figure() colors = ['green', 'orange', 'red', 'brown', 'black'] gammas = [0, 0.5, 1, 2, 5] for i, gamma in enumerate(gammas): gamma = torch.Tensor([gamma]).to(args.device) target = [] all_probs = [] for prob in np.arange(0.01, 1, 0.01): all_probs.append(torch.Tensor([prob])) target.append(torch.Tensor([1.0])) target = torch.cat(target).to(args.device) all_probs = torch.cat(all_probs).to(args.device) target = target.view(-1, 1) all_probs = all_probs.view(-1, 1) loss = FocalLoss(args, gamma=gamma, logits=False)(all_probs, target) loss = loss.cpu().numpy() plt.plot(loss, color=colors[i], label='gamma: {}'.format(str(gamma))) plt.legend() plt.show() plt.close()