├── .gitignore ├── README.md ├── config.py ├── generating_queries ├── generate_training_tuples_baseline.py ├── generate_training_tuples_refine.py └── generate_test_sets.py ├── loss └── pointnetvlad_loss.py ├── evaluate.py ├── loading_pointclouds.py ├── models └── PointNetVlad.py └── train_pointnetvlad.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.pickle 3 | log* 4 | results* 5 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PointNetVlad-Pytorch 2 | Unofficial PyTorch implementation of PointNetVlad (https://github.com/mikacuy/pointnetvlad) 3 | 4 | :warning: This repository is not maintained. Any questions regarding the approach should be directed to the authors of the original implementation: https://github.com/mikacuy/pointnetvlad 5 | 6 | I kept almost everything not related to tensorflow as the original implementation. 7 | The main differences are: 8 | * Multi-GPU support 9 | * Configuration file (config.py) 10 | * Evaluation on the eval dataset after every epochs 11 | 12 | This implementation achieved an average top 1% recall on oxford baseline of 84.81% 13 | 14 | ### Pre-Requisites 15 | * PyTorch 0.4.0 16 | * tensorboardX 17 | 18 | ### Generate pickle files 19 | ``` 20 | cd generating_queries/ 21 | 22 | # For training tuples in our baseline network 23 | python generate_training_tuples_baseline.py 24 | 25 | # For training tuples in our refined network 26 | python generate_training_tuples_refine.py 27 | 28 | # For network evaluation 29 | python generate_test_sets.py 30 | ``` 31 | 32 | ### Train 33 | ``` 34 | python train_pointnetvlad.py --dataset_folder $DATASET_FOLDER 35 | ``` 36 | 37 | ### Evaluate 38 | ``` 39 | python evaluate.py --dataset_folder $DATASET_FOLDER 40 | ``` 41 | 42 | Take a look at train_pointnetvlad.py and evaluate.py for more parameters 43 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | # GLOBAL 2 | NUM_POINTS = 4096 3 | FEATURE_OUTPUT_DIM = 256 4 | RESULTS_FOLDER = "results/" 5 | OUTPUT_FILE = "results/results.txt" 6 | 7 | LOG_DIR = 'log/' 8 | MODEL_FILENAME = "model.ckpt" 9 | 10 | DATASET_FOLDER = '../../benchmark_datasets/' 11 | 12 | # TRAIN 13 | BATCH_NUM_QUERIES = 2 14 | TRAIN_POSITIVES_PER_QUERY = 2 15 | TRAIN_NEGATIVES_PER_QUERY = 18 16 | DECAY_STEP = 200000 17 | DECAY_RATE = 0.7 18 | BASE_LEARNING_RATE = 0.000005 19 | MOMENTUM = 0.9 20 | OPTIMIZER = 'ADAM' 21 | MAX_EPOCH = 20 22 | 23 | MARGIN_1 = 0.5 24 | MARGIN_2 = 0.2 25 | 26 | BN_INIT_DECAY = 0.5 27 | BN_DECAY_DECAY_RATE = 0.5 28 | BN_DECAY_CLIP = 0.99 29 | 30 | RESUME = False 31 | 32 | TRAIN_FILE = 'generating_queries/training_queries_baseline.pickle' 33 | TEST_FILE = 'generating_queries/test_queries_baseline.pickle' 34 | 35 | # LOSS 36 | LOSS_FUNCTION = 'quadruplet' 37 | LOSS_LAZY = True 38 | TRIPLET_USE_BEST_POSITIVES = False 39 | LOSS_IGNORE_ZERO_BATCH = False 40 | 41 | # EVAL6 42 | EVAL_BATCH_SIZE = 2 43 | EVAL_POSITIVES_PER_QUERY = 4 44 | EVAL_NEGATIVES_PER_QUERY = 12 45 | 46 | EVAL_DATABASE_FILE = 'generating_queries/oxford_evaluation_database.pickle' 47 | EVAL_QUERY_FILE = 'generating_queries/oxford_evaluation_query.pickle' 48 | 49 | 50 | def cfg_str(): 51 | out_string = "" 52 | for name in globals(): 53 | if not name.startswith("__") and not name.__contains__("cfg_str"): 54 | #print(name, "=", globals()[name]) 55 | out_string = out_string + "cfg." + name + \ 56 | "=" + str(globals()[name]) + "\n" 57 | return out_string 58 | -------------------------------------------------------------------------------- /generating_queries/generate_training_tuples_baseline.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import random 4 | 5 | import numpy as np 6 | import pandas as pd 7 | from sklearn.neighbors import KDTree 8 | 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | base_path = cfg.DATASET_FOLDER 11 | 12 | runs_folder = "oxford/" 13 | filename = "pointcloud_locations_20m_10overlap.csv" 14 | pointcloud_fols = "/pointcloud_20m_10overlap/" 15 | 16 | all_folders = sorted(os.listdir(os.path.join(BASE_DIR,base_path,runs_folder))) 17 | 18 | folders = [] 19 | 20 | # All runs are used for training (both full and partial) 21 | index_list = range(len(all_folders)-1) 22 | print("Number of runs: "+str(len(index_list))) 23 | for index in index_list: 24 | folders.append(all_folders[index]) 25 | print(folders) 26 | 27 | #####For training and test data split##### 28 | x_width = 150 29 | y_width = 150 30 | p1 = [5735712.768124,620084.402381] 31 | p2 = [5735611.299219,620540.270327] 32 | p3 = [5735237.358209,620543.094379] 33 | p4 = [5734749.303802,619932.693364] 34 | p = [p1,p2,p3,p4] 35 | 36 | 37 | def check_in_test_set(northing, easting, points, x_width, y_width): 38 | in_test_set = False 39 | for point in points: 40 | if(point[0]-x_width < northing and northing < point[0]+x_width and point[1]-y_width < easting and easting < point[1]+y_width): 41 | in_test_set = True 42 | break 43 | return in_test_set 44 | ########################################## 45 | 46 | 47 | def construct_query_dict(df_centroids, filename): 48 | tree = KDTree(df_centroids[['northing','easting']]) 49 | ind_nn = tree.query_radius(df_centroids[['northing','easting']],r=10) 50 | ind_r = tree.query_radius(df_centroids[['northing','easting']], r=50) 51 | queries = {} 52 | for i in range(len(ind_nn)): 53 | query = df_centroids.iloc[i]["file"] 54 | positives = np.setdiff1d(ind_nn[i],[i]).tolist() 55 | negatives = np.setdiff1d( 56 | df_centroids.index.values.tolist(),ind_r[i]).tolist() 57 | random.shuffle(negatives) 58 | queries[i] = {"query":query, 59 | "positives":positives,"negatives":negatives} 60 | 61 | with open(filename, 'wb') as handle: 62 | pickle.dump(queries, handle, protocol=pickle.HIGHEST_PROTOCOL) 63 | 64 | print("Done ", filename) 65 | 66 | 67 | # Initialize pandas DataFrame 68 | df_train = pd.DataFrame(columns=['file','northing','easting']) 69 | df_test = pd.DataFrame(columns=['file','northing','easting']) 70 | 71 | for folder in folders: 72 | df_locations = pd.read_csv(os.path.join( 73 | base_path,runs_folder,folder,filename),sep=',') 74 | df_locations['timestamp'] = runs_folder+folder + \ 75 | pointcloud_fols+df_locations['timestamp'].astype(str)+'.bin' 76 | df_locations = df_locations.rename(columns={'timestamp':'file'}) 77 | 78 | for index, row in df_locations.iterrows(): 79 | if(check_in_test_set(row['northing'], row['easting'], p, x_width, y_width)): 80 | df_test = df_test.append(row, ignore_index=True) 81 | else: 82 | df_train = df_train.append(row, ignore_index=True) 83 | 84 | print("Number of training submaps: "+str(len(df_train['file']))) 85 | print("Number of non-disjoint test submaps: "+str(len(df_test['file']))) 86 | construct_query_dict(df_train,"training_queries_baseline.pickle") 87 | construct_query_dict(df_test,"test_queries_baseline.pickle") 88 | -------------------------------------------------------------------------------- /loss/pointnetvlad_loss.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import math 3 | import torch 4 | 5 | 6 | def best_pos_distance(query, pos_vecs): 7 | num_pos = pos_vecs.shape[1] 8 | query_copies = query.repeat(1, int(num_pos), 1) 9 | diff = ((pos_vecs - query_copies) ** 2).sum(2) 10 | min_pos, _ = diff.min(1) 11 | max_pos, _ = diff.max(1) 12 | return min_pos, max_pos 13 | 14 | 15 | def triplet_loss(q_vec, pos_vecs, neg_vecs, margin, use_min=False, lazy=False, ignore_zero_loss=False): 16 | min_pos, max_pos = best_pos_distance(q_vec, pos_vecs) 17 | 18 | # PointNetVLAD official code use min_pos, but i think max_pos should be used 19 | if use_min: 20 | positive = min_pos 21 | else: 22 | positive = max_pos 23 | 24 | num_neg = neg_vecs.shape[1] 25 | batch = q_vec.shape[0] 26 | query_copies = q_vec.repeat(1, int(num_neg), 1) 27 | positive = positive.view(-1, 1) 28 | positive = positive.repeat(1, int(num_neg)) 29 | 30 | loss = margin + positive - ((neg_vecs - query_copies) ** 2).sum(2) 31 | loss = loss.clamp(min=0.0) 32 | if lazy: 33 | triplet_loss = loss.max(1)[0] 34 | else: 35 | triplet_loss = loss.sum(1) 36 | if ignore_zero_loss: 37 | hard_triplets = torch.gt(triplet_loss, 1e-16).float() 38 | num_hard_triplets = torch.sum(hard_triplets) 39 | triplet_loss = triplet_loss.sum() / (num_hard_triplets + 1e-16) 40 | else: 41 | triplet_loss = triplet_loss.mean() 42 | return triplet_loss 43 | 44 | 45 | def triplet_loss_wrapper(q_vec, pos_vecs, neg_vecs, other_neg, m1, m2, use_min=False, lazy=False, ignore_zero_loss=False): 46 | return triplet_loss(q_vec, pos_vecs, neg_vecs, m1, use_min, lazy, ignore_zero_loss) 47 | 48 | 49 | def quadruplet_loss(q_vec, pos_vecs, neg_vecs, other_neg, m1, m2, use_min=False, lazy=False, ignore_zero_loss=False): 50 | min_pos, max_pos = best_pos_distance(q_vec, pos_vecs) 51 | 52 | # PointNetVLAD official code use min_pos, but i think max_pos should be used 53 | if use_min: 54 | positive = min_pos 55 | else: 56 | positive = max_pos 57 | 58 | num_neg = neg_vecs.shape[1] 59 | batch = q_vec.shape[0] 60 | query_copies = q_vec.repeat(1, int(num_neg), 1) 61 | positive = positive.view(-1, 1) 62 | positive = positive.repeat(1, int(num_neg)) 63 | 64 | loss = m1 + positive - ((neg_vecs - query_copies) ** 2).sum(2) 65 | loss = loss.clamp(min=0.0) 66 | if lazy: 67 | triplet_loss = loss.max(1)[0] 68 | else: 69 | triplet_loss = loss.sum(1) 70 | if ignore_zero_loss: 71 | hard_triplets = torch.gt(triplet_loss, 1e-16).float() 72 | num_hard_triplets = torch.sum(hard_triplets) 73 | triplet_loss = triplet_loss.sum() / (num_hard_triplets + 1e-16) 74 | else: 75 | triplet_loss = triplet_loss.mean() 76 | 77 | other_neg_copies = other_neg.repeat(1, int(num_neg), 1) 78 | second_loss = m2 + positive - ((neg_vecs - other_neg_copies) ** 2).sum(2) 79 | second_loss = second_loss.clamp(min=0.0) 80 | if lazy: 81 | second_loss = second_loss.max(1)[0] 82 | else: 83 | second_loss = second_loss.sum(1) 84 | 85 | if ignore_zero_loss: 86 | hard_second = torch.gt(second_loss, 1e-16).float() 87 | num_hard_second = torch.sum(hard_second) 88 | second_loss = second_loss.sum() / (num_hard_second + 1e-16) 89 | else: 90 | second_loss = second_loss.mean() 91 | 92 | total_loss = triplet_loss + second_loss 93 | return total_loss 94 | -------------------------------------------------------------------------------- /generating_queries/generate_training_tuples_refine.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import random 4 | 5 | import numpy as np 6 | import pandas as pd 7 | from sklearn.neighbors import KDTree 8 | 9 | #####For training and test data split##### 10 | x_width = 150 11 | y_width = 150 12 | 13 | # For Oxford 14 | p1 = [5735712.768124,620084.402381] 15 | p2 = [5735611.299219,620540.270327] 16 | p3 = [5735237.358209,620543.094379] 17 | p4 = [5734749.303802,619932.693364] 18 | 19 | # For University Sector 20 | p5 = [363621.292362,142864.19756] 21 | p6 = [364788.795462,143125.746609] 22 | p7 = [363597.507711,144011.414174] 23 | 24 | # For Residential Area 25 | p8 = [360895.486453,144999.915143] 26 | p9 = [362357.024536,144894.825301] 27 | p10 = [361368.907155,145209.663042] 28 | 29 | p = [p1,p2,p3,p4,p5,p6,p7,p8,p9,p10] 30 | 31 | 32 | def check_in_test_set(northing, easting, points, x_width, y_width): 33 | in_test_set = False 34 | # print(northing) 35 | for point in points: 36 | if(point[0]-x_width < northing and northing < point[0]+x_width and point[1]-y_width < easting and easting < point[1]+y_width): 37 | in_test_set = True 38 | break 39 | return in_test_set 40 | ########################################## 41 | 42 | 43 | def construct_query_dict(df_centroids, filename): 44 | tree = KDTree(df_centroids[['northing','easting']]) 45 | ind_nn = tree.query_radius(df_centroids[['northing','easting']],r=12.5) 46 | ind_r = tree.query_radius(df_centroids[['northing','easting']], r=50) 47 | queries = {} 48 | print(len(ind_nn)) 49 | for i in range(len(ind_nn)): 50 | query = df_centroids.iloc[i]["file"] 51 | positives = np.setdiff1d(ind_nn[i],[i]).tolist() 52 | negatives = np.setdiff1d( 53 | df_centroids.index.values.tolist(),ind_r[i]).tolist() 54 | random.shuffle(negatives) 55 | queries[i] = {"query":query, 56 | "positives":positives,"negatives":negatives} 57 | 58 | with open(filename, 'wb') as handle: 59 | pickle.dump(queries, handle, protocol=pickle.HIGHEST_PROTOCOL) 60 | 61 | print("Done ", filename) 62 | 63 | 64 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 65 | base_path = cfg.DATASET_FOLDER 66 | runs_folder = "inhouse_datasets/" 67 | filename = "pointcloud_centroids_10.csv" 68 | pointcloud_fols = "/pointcloud_25m_10/" 69 | 70 | all_folders = sorted(os.listdir(os.path.join(BASE_DIR,base_path,runs_folder))) 71 | 72 | folders = [] 73 | index_list = range(5,15) 74 | for index in index_list: 75 | folders.append(all_folders[index]) 76 | 77 | print(folders) 78 | 79 | # Initialize pandas DataFrame 80 | df_train = pd.DataFrame(columns=['file','northing','easting']) 81 | 82 | for folder in folders: 83 | df_locations = pd.read_csv(os.path.join( 84 | base_path, runs_folder, folder,filename),sep=',') 85 | df_locations['timestamp'] = runs_folder+folder + \ 86 | pointcloud_fols+df_locations['timestamp'].astype(str)+'.bin' 87 | df_locations = df_locations.rename(columns={'timestamp':'file'}) 88 | for index, row in df_locations.iterrows(): 89 | if(check_in_test_set(row['northing'], row['easting'], p, x_width, y_width)): 90 | continue 91 | else: 92 | df_train = df_train.append(row, ignore_index=True) 93 | 94 | print(len(df_train['file'])) 95 | 96 | 97 | # Combine with Oxford data 98 | runs_folder = "oxford/" 99 | filename = "pointcloud_locations_20m_10overlap.csv" 100 | pointcloud_fols = "/pointcloud_20m_10overlap/" 101 | 102 | all_folders = sorted(os.listdir(os.path.join(BASE_DIR,base_path,runs_folder))) 103 | 104 | folders = [] 105 | index_list = range(len(all_folders)-1) 106 | for index in index_list: 107 | folders.append(all_folders[index]) 108 | 109 | print(folders) 110 | 111 | for folder in folders: 112 | df_locations = pd.read_csv(os.path.join( 113 | base_path,runs_folder,folder,filename),sep=',') 114 | df_locations['timestamp'] = runs_folder+folder + \ 115 | pointcloud_fols+df_locations['timestamp'].astype(str)+'.bin' 116 | df_locations = df_locations.rename(columns={'timestamp':'file'}) 117 | for index, row in df_locations.iterrows(): 118 | if(check_in_test_set(row['northing'], row['easting'], p, x_width, y_width)): 119 | continue 120 | else: 121 | df_train = df_train.append(row, ignore_index=True) 122 | 123 | print("Number of training submaps: "+str(len(df_train['file']))) 124 | construct_query_dict(df_train,"training_queries_refine.pickle") 125 | -------------------------------------------------------------------------------- /generating_queries/generate_test_sets.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import random 4 | 5 | import numpy as np 6 | import pandas as pd 7 | from sklearn.neighbors import KDTree 8 | 9 | import config as cfg 10 | 11 | #####For training and test data split##### 12 | x_width = 150 13 | y_width = 150 14 | 15 | # For Oxford 16 | p1 = [5735712.768124,620084.402381] 17 | p2 = [5735611.299219,620540.270327] 18 | p3 = [5735237.358209,620543.094379] 19 | p4 = [5734749.303802,619932.693364] 20 | 21 | # For University Sector 22 | p5 = [363621.292362,142864.19756] 23 | p6 = [364788.795462,143125.746609] 24 | p7 = [363597.507711,144011.414174] 25 | 26 | # For Residential Area 27 | p8 = [360895.486453,144999.915143] 28 | p9 = [362357.024536,144894.825301] 29 | p10 = [361368.907155,145209.663042] 30 | 31 | p_dict = {"oxford":[p1,p2,p3,p4], "university":[ 32 | p5,p6,p7], "residential": [p8,p9,p10], "business":[]} 33 | 34 | 35 | def check_in_test_set(northing, easting, points, x_width, y_width): 36 | in_test_set = False 37 | for point in points: 38 | if(point[0]-x_width < northing and northing < point[0]+x_width and point[1]-y_width < easting and easting < point[1]+y_width): 39 | in_test_set = True 40 | break 41 | return in_test_set 42 | ########################################## 43 | 44 | 45 | def output_to_file(output, filename): 46 | with open(filename, 'wb') as handle: 47 | pickle.dump(output, handle, protocol=pickle.HIGHEST_PROTOCOL) 48 | print("Done ", filename) 49 | 50 | 51 | def construct_query_and_database_sets(base_path, runs_folder, folders, pointcloud_fols, filename, p, output_name): 52 | database_trees = [] 53 | test_trees = [] 54 | for folder in folders: 55 | print(folder) 56 | df_database = pd.DataFrame(columns=['file','northing','easting']) 57 | df_test = pd.DataFrame(columns=['file','northing','easting']) 58 | 59 | df_locations = pd.read_csv(os.path.join( 60 | base_path,runs_folder,folder,filename),sep=',') 61 | # df_locations['timestamp']=runs_folder+folder+pointcloud_fols+df_locations['timestamp'].astype(str)+'.bin' 62 | # df_locations=df_locations.rename(columns={'timestamp':'file'}) 63 | for index, row in df_locations.iterrows(): 64 | # entire business district is in the test set 65 | if(output_name == "business"): 66 | df_test = df_test.append(row, ignore_index=True) 67 | elif(check_in_test_set(row['northing'], row['easting'], p, x_width, y_width)): 68 | df_test = df_test.append(row, ignore_index=True) 69 | df_database = df_database.append(row, ignore_index=True) 70 | 71 | database_tree = KDTree(df_database[['northing','easting']]) 72 | test_tree = KDTree(df_test[['northing','easting']]) 73 | database_trees.append(database_tree) 74 | test_trees.append(test_tree) 75 | 76 | test_sets = [] 77 | database_sets = [] 78 | for folder in folders: 79 | database = {} 80 | test = {} 81 | df_locations = pd.read_csv(os.path.join( 82 | base_path,runs_folder,folder,filename),sep=',') 83 | df_locations['timestamp'] = runs_folder+folder + \ 84 | pointcloud_fols+df_locations['timestamp'].astype(str)+'.bin' 85 | df_locations = df_locations.rename(columns={'timestamp':'file'}) 86 | for index,row in df_locations.iterrows(): 87 | # entire business district is in the test set 88 | if(output_name == "business"): 89 | test[len(test.keys())] = { 90 | 'query':row['file'],'northing':row['northing'],'easting':row['easting']} 91 | elif(check_in_test_set(row['northing'], row['easting'], p, x_width, y_width)): 92 | test[len(test.keys())] = { 93 | 'query':row['file'],'northing':row['northing'],'easting':row['easting']} 94 | database[len(database.keys())] = { 95 | 'query':row['file'],'northing':row['northing'],'easting':row['easting']} 96 | database_sets.append(database) 97 | test_sets.append(test) 98 | 99 | for i in range(len(database_sets)): 100 | tree = database_trees[i] 101 | for j in range(len(test_sets)): 102 | if(i == j): 103 | continue 104 | for key in range(len(test_sets[j].keys())): 105 | coor = np.array( 106 | [[test_sets[j][key]["northing"],test_sets[j][key]["easting"]]]) 107 | index = tree.query_radius(coor, r=25) 108 | # indices of the positive matches in database i of each query (key) in test set j 109 | test_sets[j][key][i] = index[0].tolist() 110 | 111 | output_to_file(database_sets, output_name+'_evaluation_database.pickle') 112 | output_to_file(test_sets, output_name+'_evaluation_query.pickle') 113 | 114 | 115 | # Building database and query files for evaluation 116 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 117 | base_path = cfg.DATASET_FOLDER 118 | 119 | # For Oxford 120 | folders = [] 121 | runs_folder = "oxford/" 122 | all_folders = sorted(os.listdir(os.path.join(BASE_DIR,base_path,runs_folder))) 123 | index_list = [5,6,7,9,10,11,12,13,14,15,16,17,18,19,22,24,31,32,33,38,39,43,44] 124 | print(len(index_list)) 125 | for index in index_list: 126 | folders.append(all_folders[index]) 127 | 128 | print(folders) 129 | construct_query_and_database_sets(base_path, runs_folder, folders, "/pointcloud_20m/", 130 | "pointcloud_locations_20m.csv", p_dict["oxford"], "oxford") 131 | 132 | # For University Sector 133 | folders = [] 134 | runs_folder = "inhouse_datasets/" 135 | all_folders = sorted(os.listdir(os.path.join(BASE_DIR,base_path,runs_folder))) 136 | uni_index = range(10,15) 137 | for index in uni_index: 138 | folders.append(all_folders[index]) 139 | 140 | print(folders) 141 | construct_query_and_database_sets(base_path, runs_folder, folders, "/pointcloud_25m_25/", 142 | "pointcloud_centroids_25.csv", p_dict["university"], "university") 143 | 144 | # For Residential Area 145 | folders = [] 146 | runs_folder = "inhouse_datasets/" 147 | all_folders = sorted(os.listdir(os.path.join(BASE_DIR,base_path,runs_folder))) 148 | res_index = range(5,10) 149 | for index in res_index: 150 | folders.append(all_folders[index]) 151 | 152 | print(folders) 153 | construct_query_and_database_sets(base_path, runs_folder, folders, "/pointcloud_25m_25/", 154 | "pointcloud_centroids_25.csv", p_dict["residential"], "residential") 155 | 156 | # For Business District 157 | folders = [] 158 | runs_folder = "inhouse_datasets/" 159 | all_folders = sorted(os.listdir(os.path.join(BASE_DIR,base_path,runs_folder))) 160 | bus_index = range(5) 161 | for index in bus_index: 162 | folders.append(all_folders[index]) 163 | 164 | print(folders) 165 | construct_query_and_database_sets(base_path, runs_folder, folders, "/pointcloud_25m_25/", 166 | "pointcloud_centroids_25.csv", p_dict["business"], "business") 167 | -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | import numpy as np 4 | import socket 5 | import importlib 6 | import os 7 | os.environ["CUDA_VISIBLE_DEVICES"] = "3" 8 | import sys 9 | import torch 10 | import torch.nn as nn 11 | from torch.autograd import Variable 12 | from torch.backends import cudnn 13 | 14 | from sklearn.neighbors import NearestNeighbors 15 | from sklearn.neighbors import KDTree 16 | 17 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 18 | sys.path.append(BASE_DIR) 19 | 20 | from loading_pointclouds import * 21 | import models.PointNetVlad as PNV 22 | from tensorboardX import SummaryWriter 23 | import loss.pointnetvlad_loss 24 | 25 | import config as cfg 26 | 27 | cudnn.enabled = True 28 | 29 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 30 | 31 | 32 | def evaluate(): 33 | model = PNV.PointNetVlad(global_feat=True, feature_transform=True, max_pool=False, 34 | output_dim=cfg.FEATURE_OUTPUT_DIM, num_points=cfg.NUM_POINTS) 35 | model = model.to(device) 36 | 37 | resume_filename = cfg.LOG_DIR + "checkpoint.pth.tar" 38 | print("Resuming From ", resume_filename) 39 | checkpoint = torch.load(resume_filename) 40 | saved_state_dict = checkpoint['state_dict'] 41 | model.load_state_dict(saved_state_dict) 42 | 43 | model = nn.DataParallel(model) 44 | 45 | print(evaluate_model(model)) 46 | 47 | 48 | def evaluate_model(model): 49 | DATABASE_SETS = get_sets_dict(cfg.EVAL_DATABASE_FILE) 50 | QUERY_SETS = get_sets_dict(cfg.EVAL_QUERY_FILE) 51 | 52 | if not os.path.exists(cfg.RESULTS_FOLDER): 53 | os.mkdir(cfg.RESULTS_FOLDER) 54 | 55 | recall = np.zeros(25) 56 | count = 0 57 | similarity = [] 58 | one_percent_recall = [] 59 | 60 | DATABASE_VECTORS = [] 61 | QUERY_VECTORS = [] 62 | 63 | for i in range(len(DATABASE_SETS)): 64 | DATABASE_VECTORS.append(get_latent_vectors(model, DATABASE_SETS[i])) 65 | 66 | for j in range(len(QUERY_SETS)): 67 | QUERY_VECTORS.append(get_latent_vectors(model, QUERY_SETS[j])) 68 | 69 | for m in range(len(QUERY_SETS)): 70 | for n in range(len(QUERY_SETS)): 71 | if (m == n): 72 | continue 73 | pair_recall, pair_similarity, pair_opr = get_recall( 74 | m, n, DATABASE_VECTORS, QUERY_VECTORS, QUERY_SETS) 75 | recall += np.array(pair_recall) 76 | count += 1 77 | one_percent_recall.append(pair_opr) 78 | for x in pair_similarity: 79 | similarity.append(x) 80 | 81 | print() 82 | ave_recall = recall / count 83 | # print(ave_recall) 84 | 85 | # print(similarity) 86 | average_similarity = np.mean(similarity) 87 | # print(average_similarity) 88 | 89 | ave_one_percent_recall = np.mean(one_percent_recall) 90 | # print(ave_one_percent_recall) 91 | 92 | with open(cfg.OUTPUT_FILE, "w") as output: 93 | output.write("Average Recall @N:\n") 94 | output.write(str(ave_recall)) 95 | output.write("\n\n") 96 | output.write("Average Similarity:\n") 97 | output.write(str(average_similarity)) 98 | output.write("\n\n") 99 | output.write("Average Top 1% Recall:\n") 100 | output.write(str(ave_one_percent_recall)) 101 | 102 | return ave_one_percent_recall 103 | 104 | 105 | def get_latent_vectors(model, dict_to_process): 106 | 107 | model.eval() 108 | is_training = False 109 | train_file_idxs = np.arange(0, len(dict_to_process.keys())) 110 | 111 | batch_num = cfg.EVAL_BATCH_SIZE * \ 112 | (1 + cfg.EVAL_POSITIVES_PER_QUERY + cfg.EVAL_NEGATIVES_PER_QUERY) 113 | q_output = [] 114 | for q_index in range(len(train_file_idxs)//batch_num): 115 | file_indices = train_file_idxs[q_index * 116 | batch_num:(q_index+1)*(batch_num)] 117 | file_names = [] 118 | for index in file_indices: 119 | file_names.append(dict_to_process[index]["query"]) 120 | queries = load_pc_files(file_names) 121 | 122 | with torch.no_grad(): 123 | feed_tensor = torch.from_numpy(queries).float() 124 | feed_tensor = feed_tensor.unsqueeze(1) 125 | feed_tensor = feed_tensor.to(device) 126 | out = model(feed_tensor) 127 | 128 | out = out.detach().cpu().numpy() 129 | out = np.squeeze(out) 130 | 131 | #out = np.vstack((o1, o2, o3, o4)) 132 | q_output.append(out) 133 | 134 | q_output = np.array(q_output) 135 | if(len(q_output) != 0): 136 | q_output = q_output.reshape(-1, q_output.shape[-1]) 137 | 138 | # handle edge case 139 | index_edge = len(train_file_idxs) // batch_num * batch_num 140 | if index_edge < len(dict_to_process.keys()): 141 | file_indices = train_file_idxs[index_edge:len(dict_to_process.keys())] 142 | file_names = [] 143 | for index in file_indices: 144 | file_names.append(dict_to_process[index]["query"]) 145 | queries = load_pc_files(file_names) 146 | 147 | with torch.no_grad(): 148 | feed_tensor = torch.from_numpy(queries).float() 149 | feed_tensor = feed_tensor.unsqueeze(1) 150 | feed_tensor = feed_tensor.to(device) 151 | o1 = model(feed_tensor) 152 | 153 | output = o1.detach().cpu().numpy() 154 | output = np.squeeze(output) 155 | if (q_output.shape[0] != 0): 156 | q_output = np.vstack((q_output, output)) 157 | else: 158 | q_output = output 159 | 160 | model.train() 161 | # print(q_output.shape) 162 | return q_output 163 | 164 | 165 | def get_recall(m, n, DATABASE_VECTORS, QUERY_VECTORS, QUERY_SETS): 166 | 167 | database_output = DATABASE_VECTORS[m] 168 | queries_output = QUERY_VECTORS[n] 169 | 170 | # print(len(queries_output)) 171 | database_nbrs = KDTree(database_output) 172 | 173 | num_neighbors = 25 174 | recall = [0] * num_neighbors 175 | 176 | top1_similarity_score = [] 177 | one_percent_retrieved = 0 178 | threshold = max(int(round(len(database_output)/100.0)), 1) 179 | 180 | num_evaluated = 0 181 | for i in range(len(queries_output)): 182 | true_neighbors = QUERY_SETS[n][i][m] 183 | if(len(true_neighbors) == 0): 184 | continue 185 | num_evaluated += 1 186 | distances, indices = database_nbrs.query( 187 | np.array([queries_output[i]]),k=num_neighbors) 188 | for j in range(len(indices[0])): 189 | if indices[0][j] in true_neighbors: 190 | if(j == 0): 191 | similarity = np.dot( 192 | queries_output[i], database_output[indices[0][j]]) 193 | top1_similarity_score.append(similarity) 194 | recall[j] += 1 195 | break 196 | 197 | if len(list(set(indices[0][0:threshold]).intersection(set(true_neighbors)))) > 0: 198 | one_percent_retrieved += 1 199 | 200 | one_percent_recall = (one_percent_retrieved/float(num_evaluated))*100 201 | recall = (np.cumsum(recall)/float(num_evaluated))*100 202 | # print(recall) 203 | # print(np.mean(top1_similarity_score)) 204 | # print(one_percent_recall) 205 | return recall, top1_similarity_score, one_percent_recall 206 | 207 | 208 | if __name__ == "__main__": 209 | # params 210 | parser = argparse.ArgumentParser() 211 | parser.add_argument('--positives_per_query', type=int, default=4, 212 | help='Number of potential positives in each training tuple [default: 2]') 213 | parser.add_argument('--negatives_per_query', type=int, default=12, 214 | help='Number of definite negatives in each training tuple [default: 20]') 215 | parser.add_argument('--eval_batch_size', type=int, default=12, 216 | help='Batch Size during training [default: 1]') 217 | parser.add_argument('--dimension', type=int, default=256) 218 | parser.add_argument('--decay_step', type=int, default=200000, 219 | help='Decay step for lr decay [default: 200000]') 220 | parser.add_argument('--decay_rate', type=float, default=0.7, 221 | help='Decay rate for lr decay [default: 0.8]') 222 | parser.add_argument('--results_dir', default='results/', 223 | help='results dir [default: results]') 224 | parser.add_argument('--dataset_folder', default='../../dataset/', 225 | help='PointNetVlad Dataset Folder') 226 | FLAGS = parser.parse_args() 227 | 228 | #BATCH_SIZE = FLAGS.batch_size 229 | #cfg.EVAL_BATCH_SIZE = FLAGS.eval_batch_size 230 | cfg.NUM_POINTS = 4096 231 | cfg.FEATURE_OUTPUT_DIM = 256 232 | cfg.EVAL_POSITIVES_PER_QUERY = FLAGS.positives_per_query 233 | cfg.EVAL_NEGATIVES_PER_QUERY = FLAGS.negatives_per_query 234 | cfg.DECAY_STEP = FLAGS.decay_step 235 | cfg.DECAY_RATE = FLAGS.decay_rate 236 | 237 | cfg.RESULTS_FOLDER = FLAGS.results_dir 238 | 239 | cfg.EVAL_DATABASE_FILE = 'generating_queries/oxford_evaluation_database.pickle' 240 | cfg.EVAL_QUERY_FILE = 'generating_queries/oxford_evaluation_query.pickle' 241 | 242 | cfg.LOG_DIR = 'log/' 243 | cfg.OUTPUT_FILE = cfg.RESULTS_FOLDER + 'results.txt' 244 | cfg.MODEL_FILENAME = "model.ckpt" 245 | 246 | cfg.DATASET_FOLDER = FLAGS.dataset_folder 247 | 248 | evaluate() 249 | -------------------------------------------------------------------------------- /loading_pointclouds.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import numpy as np 4 | import random 5 | import config as cfg 6 | 7 | def get_queries_dict(filename): 8 | # key:{'query':file,'positives':[files],'negatives:[files], 'neighbors':[keys]} 9 | with open(filename, 'rb') as handle: 10 | queries = pickle.load(handle) 11 | print("Queries Loaded.") 12 | return queries 13 | 14 | 15 | def get_sets_dict(filename): 16 | #[key_dataset:{key_pointcloud:{'query':file,'northing':value,'easting':value}},key_dataset:{key_pointcloud:{'query':file,'northing':value,'easting':value}}, ...} 17 | with open(filename, 'rb') as handle: 18 | trajectories = pickle.load(handle) 19 | print("Trajectories Loaded.") 20 | return trajectories 21 | 22 | 23 | def load_pc_file(filename): 24 | # returns Nx3 matrix 25 | pc = np.fromfile(os.path.join(cfg.DATASET_FOLDER, filename), dtype=np.float64) 26 | 27 | if(pc.shape[0] != 4096*3): 28 | print("Error in pointcloud shape") 29 | return np.array([]) 30 | 31 | pc = np.reshape(pc,(pc.shape[0]//3, 3)) 32 | return pc 33 | 34 | 35 | def load_pc_files(filenames): 36 | pcs = [] 37 | for filename in filenames: 38 | # print(filename) 39 | pc = load_pc_file(filename) 40 | if(pc.shape[0] != 4096): 41 | continue 42 | pcs.append(pc) 43 | pcs = np.array(pcs) 44 | return pcs 45 | 46 | 47 | def rotate_point_cloud(batch_data): 48 | """ Randomly rotate the point clouds to augument the dataset 49 | rotation is per shape based along up direction 50 | Input: 51 | BxNx3 array, original batch of point clouds 52 | Return: 53 | BxNx3 array, rotated batch of point clouds 54 | """ 55 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 56 | for k in range(batch_data.shape[0]): 57 | #rotation_angle = np.random.uniform() * 2 * np.pi 58 | #-90 to 90 59 | rotation_angle = (np.random.uniform()*np.pi) - np.pi/2.0 60 | cosval = np.cos(rotation_angle) 61 | sinval = np.sin(rotation_angle) 62 | rotation_matrix = np.array([[cosval, -sinval, 0], 63 | [sinval, cosval, 0], 64 | [0, 0, 1]]) 65 | shape_pc = batch_data[k, ...] 66 | rotated_data[k, ...] = np.dot( 67 | shape_pc.reshape((-1, 3)), rotation_matrix) 68 | return rotated_data 69 | 70 | 71 | def jitter_point_cloud(batch_data, sigma=0.005, clip=0.05): 72 | """ Randomly jitter points. jittering is per point. 73 | Input: 74 | BxNx3 array, original batch of point clouds 75 | Return: 76 | BxNx3 array, jittered batch of point clouds 77 | """ 78 | B, N, C = batch_data.shape 79 | assert(clip > 0) 80 | jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip) 81 | jittered_data += batch_data 82 | return jittered_data 83 | 84 | 85 | def get_query_tuple(dict_value, num_pos, num_neg, QUERY_DICT, hard_neg=[], other_neg=False): 86 | # get query tuple for dictionary entry 87 | # return list [query,positives,negatives] 88 | 89 | query = load_pc_file(dict_value["query"]) # Nx3 90 | 91 | random.shuffle(dict_value["positives"]) 92 | pos_files = [] 93 | 94 | for i in range(num_pos): 95 | pos_files.append(QUERY_DICT[dict_value["positives"][i]]["query"]) 96 | #positives= load_pc_files(dict_value["positives"][0:num_pos]) 97 | positives = load_pc_files(pos_files) 98 | 99 | neg_files = [] 100 | neg_indices = [] 101 | if(len(hard_neg) == 0): 102 | random.shuffle(dict_value["negatives"]) 103 | for i in range(num_neg): 104 | neg_files.append(QUERY_DICT[dict_value["negatives"][i]]["query"]) 105 | neg_indices.append(dict_value["negatives"][i]) 106 | 107 | else: 108 | random.shuffle(dict_value["negatives"]) 109 | for i in hard_neg: 110 | neg_files.append(QUERY_DICT[i]["query"]) 111 | neg_indices.append(i) 112 | j = 0 113 | while(len(neg_files) < num_neg): 114 | 115 | if not dict_value["negatives"][j] in hard_neg: 116 | neg_files.append( 117 | QUERY_DICT[dict_value["negatives"][j]]["query"]) 118 | neg_indices.append(dict_value["negatives"][j]) 119 | j += 1 120 | 121 | negatives = load_pc_files(neg_files) 122 | 123 | if other_neg is False: 124 | return [query, positives, negatives] 125 | # For Quadruplet Loss 126 | else: 127 | # get neighbors of negatives and query 128 | neighbors = [] 129 | for pos in dict_value["positives"]: 130 | neighbors.append(pos) 131 | for neg in neg_indices: 132 | for pos in QUERY_DICT[neg]["positives"]: 133 | neighbors.append(pos) 134 | possible_negs = list(set(QUERY_DICT.keys())-set(neighbors)) 135 | random.shuffle(possible_negs) 136 | 137 | if(len(possible_negs) == 0): 138 | return [query, positives, negatives, np.array([])] 139 | 140 | neg2 = load_pc_file(QUERY_DICT[possible_negs[0]]["query"]) 141 | 142 | return [query, positives, negatives, neg2] 143 | 144 | 145 | def get_rotated_tuple(dict_value, num_pos, num_neg, QUERY_DICT, hard_neg=[], other_neg=False): 146 | query = load_pc_file(dict_value["query"]) # Nx3 147 | q_rot = rotate_point_cloud(np.expand_dims(query, axis=0)) 148 | q_rot = np.squeeze(q_rot) 149 | 150 | random.shuffle(dict_value["positives"]) 151 | pos_files = [] 152 | for i in range(num_pos): 153 | pos_files.append(QUERY_DICT[dict_value["positives"][i]]["query"]) 154 | #positives= load_pc_files(dict_value["positives"][0:num_pos]) 155 | positives = load_pc_files(pos_files) 156 | p_rot = rotate_point_cloud(positives) 157 | 158 | neg_files = [] 159 | neg_indices = [] 160 | if(len(hard_neg) == 0): 161 | random.shuffle(dict_value["negatives"]) 162 | for i in range(num_neg): 163 | neg_files.append(QUERY_DICT[dict_value["negatives"][i]]["query"]) 164 | neg_indices.append(dict_value["negatives"][i]) 165 | else: 166 | random.shuffle(dict_value["negatives"]) 167 | for i in hard_neg: 168 | neg_files.append(QUERY_DICT[i]["query"]) 169 | neg_indices.append(i) 170 | j = 0 171 | while(len(neg_files) < num_neg): 172 | if not dict_value["negatives"][j] in hard_neg: 173 | neg_files.append( 174 | QUERY_DICT[dict_value["negatives"][j]]["query"]) 175 | neg_indices.append(dict_value["negatives"][j]) 176 | j += 1 177 | negatives = load_pc_files(neg_files) 178 | n_rot = rotate_point_cloud(negatives) 179 | 180 | if other_neg is False: 181 | return [q_rot, p_rot, n_rot] 182 | 183 | # For Quadruplet Loss 184 | else: 185 | # get neighbors of negatives and query 186 | neighbors = [] 187 | for pos in dict_value["positives"]: 188 | neighbors.append(pos) 189 | for neg in neg_indices: 190 | for pos in QUERY_DICT[neg]["positives"]: 191 | neighbors.append(pos) 192 | possible_negs = list(set(QUERY_DICT.keys())-set(neighbors)) 193 | random.shuffle(possible_negs) 194 | 195 | if(len(possible_negs) == 0): 196 | return [q_jit, p_jit, n_jit, np.array([])] 197 | 198 | neg2 = load_pc_file(QUERY_DICT[possible_negs[0]]["query"]) 199 | n2_rot = rotate_point_cloud(np.expand_dims(neg2, axis=0)) 200 | n2_rot = np.squeeze(n2_rot) 201 | 202 | return [q_rot, p_rot, n_rot, n2_rot] 203 | 204 | 205 | def get_jittered_tuple(dict_value, num_pos, num_neg, QUERY_DICT, hard_neg=[], other_neg=False): 206 | query = load_pc_file(dict_value["query"]) # Nx3 207 | #q_rot= rotate_point_cloud(np.expand_dims(query, axis=0)) 208 | q_jit = jitter_point_cloud(np.expand_dims(query, axis=0)) 209 | q_jit = np.squeeze(q_jit) 210 | 211 | random.shuffle(dict_value["positives"]) 212 | pos_files = [] 213 | for i in range(num_pos): 214 | pos_files.append(QUERY_DICT[dict_value["positives"][i]]["query"]) 215 | #positives= load_pc_files(dict_value["positives"][0:num_pos]) 216 | positives = load_pc_files(pos_files) 217 | p_jit = jitter_point_cloud(positives) 218 | 219 | neg_files = [] 220 | neg_indices = [] 221 | if(len(hard_neg) == 0): 222 | random.shuffle(dict_value["negatives"]) 223 | for i in range(num_neg): 224 | neg_files.append(QUERY_DICT[dict_value["negatives"][i]]["query"]) 225 | neg_indices.append(dict_value["negatives"][i]) 226 | else: 227 | random.shuffle(dict_value["negatives"]) 228 | for i in hard_neg: 229 | neg_files.append(QUERY_DICT[i]["query"]) 230 | neg_indices.append(i) 231 | j = 0 232 | while(len(neg_files) < num_neg): 233 | if not dict_value["negatives"][j] in hard_neg: 234 | neg_files.append( 235 | QUERY_DICT[dict_value["negatives"][j]]["query"]) 236 | neg_indices.append(dict_value["negatives"][j]) 237 | j += 1 238 | negatives = load_pc_files(neg_files) 239 | n_jit = jitter_point_cloud(negatives) 240 | 241 | if other_neg is False: 242 | return [q_jit, p_jit, n_jit] 243 | 244 | # For Quadruplet Loss 245 | else: 246 | # get neighbors of negatives and query 247 | neighbors = [] 248 | for pos in dict_value["positives"]: 249 | neighbors.append(pos) 250 | for neg in neg_indices: 251 | for pos in QUERY_DICT[neg]["positives"]: 252 | neighbors.append(pos) 253 | possible_negs = list(set(QUERY_DICT.keys())-set(neighbors)) 254 | random.shuffle(possible_negs) 255 | 256 | if(len(possible_negs) == 0): 257 | return [q_jit, p_jit, n_jit, np.array([])] 258 | 259 | neg2 = load_pc_file(QUERY_DICT[possible_negs[0]]["query"]) 260 | n2_jit = jitter_point_cloud(np.expand_dims(neg2, axis=0)) 261 | n2_jit = np.squeeze(n2_jit) 262 | 263 | return [q_jit, p_jit, n_jit, n2_jit] 264 | -------------------------------------------------------------------------------- /models/PointNetVlad.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.parallel 5 | import torch.utils.data 6 | from torch.autograd import Variable 7 | import numpy as np 8 | import torch.nn.functional as F 9 | import math 10 | 11 | 12 | class NetVLADLoupe(nn.Module): 13 | def __init__(self, feature_size, max_samples, cluster_size, output_dim, 14 | gating=True, add_batch_norm=True, is_training=True): 15 | super(NetVLADLoupe, self).__init__() 16 | self.feature_size = feature_size 17 | self.max_samples = max_samples 18 | self.output_dim = output_dim 19 | self.is_training = is_training 20 | self.gating = gating 21 | self.add_batch_norm = add_batch_norm 22 | self.cluster_size = cluster_size 23 | self.softmax = nn.Softmax(dim=-1) 24 | self.cluster_weights = nn.Parameter(torch.randn( 25 | feature_size, cluster_size) * 1 / math.sqrt(feature_size)) 26 | self.cluster_weights2 = nn.Parameter(torch.randn( 27 | 1, feature_size, cluster_size) * 1 / math.sqrt(feature_size)) 28 | self.hidden1_weights = nn.Parameter( 29 | torch.randn(cluster_size * feature_size, output_dim) * 1 / math.sqrt(feature_size)) 30 | 31 | if add_batch_norm: 32 | self.cluster_biases = None 33 | self.bn1 = nn.BatchNorm1d(cluster_size) 34 | else: 35 | self.cluster_biases = nn.Parameter(torch.randn( 36 | cluster_size) * 1 / math.sqrt(feature_size)) 37 | self.bn1 = None 38 | 39 | self.bn2 = nn.BatchNorm1d(output_dim) 40 | 41 | if gating: 42 | self.context_gating = GatingContext( 43 | output_dim, add_batch_norm=add_batch_norm) 44 | 45 | def forward(self, x): 46 | x = x.transpose(1, 3).contiguous() 47 | x = x.view((-1, self.max_samples, self.feature_size)) 48 | activation = torch.matmul(x, self.cluster_weights) 49 | if self.add_batch_norm: 50 | # activation = activation.transpose(1,2).contiguous() 51 | activation = activation.view(-1, self.cluster_size) 52 | activation = self.bn1(activation) 53 | activation = activation.view(-1, 54 | self.max_samples, self.cluster_size) 55 | # activation = activation.transpose(1,2).contiguous() 56 | else: 57 | activation = activation + self.cluster_biases 58 | activation = self.softmax(activation) 59 | activation = activation.view((-1, self.max_samples, self.cluster_size)) 60 | 61 | a_sum = activation.sum(-2, keepdim=True) 62 | a = a_sum * self.cluster_weights2 63 | 64 | activation = torch.transpose(activation, 2, 1) 65 | x = x.view((-1, self.max_samples, self.feature_size)) 66 | vlad = torch.matmul(activation, x) 67 | vlad = torch.transpose(vlad, 2, 1) 68 | vlad = vlad - a 69 | 70 | vlad = F.normalize(vlad, dim=1, p=2) 71 | vlad = vlad.view((-1, self.cluster_size * self.feature_size)) 72 | vlad = F.normalize(vlad, dim=1, p=2) 73 | 74 | vlad = torch.matmul(vlad, self.hidden1_weights) 75 | 76 | vlad = self.bn2(vlad) 77 | 78 | if self.gating: 79 | vlad = self.context_gating(vlad) 80 | 81 | return vlad 82 | 83 | 84 | class GatingContext(nn.Module): 85 | def __init__(self, dim, add_batch_norm=True): 86 | super(GatingContext, self).__init__() 87 | self.dim = dim 88 | self.add_batch_norm = add_batch_norm 89 | self.gating_weights = nn.Parameter( 90 | torch.randn(dim, dim) * 1 / math.sqrt(dim)) 91 | self.sigmoid = nn.Sigmoid() 92 | 93 | if add_batch_norm: 94 | self.gating_biases = None 95 | self.bn1 = nn.BatchNorm1d(dim) 96 | else: 97 | self.gating_biases = nn.Parameter( 98 | torch.randn(dim) * 1 / math.sqrt(dim)) 99 | self.bn1 = None 100 | 101 | def forward(self, x): 102 | gates = torch.matmul(x, self.gating_weights) 103 | 104 | if self.add_batch_norm: 105 | gates = self.bn1(gates) 106 | else: 107 | gates = gates + self.gating_biases 108 | 109 | gates = self.sigmoid(gates) 110 | 111 | activation = x * gates 112 | 113 | return activation 114 | 115 | 116 | class Flatten(nn.Module): 117 | def __init__(self): 118 | nn.Module.__init__(self) 119 | 120 | def forward(self, input): 121 | return input.view(input.size(0), -1) 122 | 123 | 124 | class STN3d(nn.Module): 125 | def __init__(self, num_points=2500, k=3, use_bn=True): 126 | super(STN3d, self).__init__() 127 | self.k = k 128 | self.kernel_size = 3 if k == 3 else 1 129 | self.channels = 1 if k == 3 else k 130 | self.num_points = num_points 131 | self.use_bn = use_bn 132 | self.conv1 = torch.nn.Conv2d(self.channels, 64, (1, self.kernel_size)) 133 | self.conv2 = torch.nn.Conv2d(64, 128, (1,1)) 134 | self.conv3 = torch.nn.Conv2d(128, 1024, (1,1)) 135 | self.mp1 = torch.nn.MaxPool2d((num_points, 1), 1) 136 | self.fc1 = nn.Linear(1024, 512) 137 | self.fc2 = nn.Linear(512, 256) 138 | self.fc3 = nn.Linear(256, k*k) 139 | self.fc3.weight.data.zero_() 140 | self.fc3.bias.data.zero_() 141 | self.relu = nn.ReLU() 142 | 143 | if use_bn: 144 | self.bn1 = nn.BatchNorm2d(64) 145 | self.bn2 = nn.BatchNorm2d(128) 146 | self.bn3 = nn.BatchNorm2d(1024) 147 | self.bn4 = nn.BatchNorm1d(512) 148 | self.bn5 = nn.BatchNorm1d(256) 149 | 150 | def forward(self, x): 151 | batchsize = x.size()[0] 152 | if self.use_bn: 153 | x = F.relu(self.bn1(self.conv1(x))) 154 | x = F.relu(self.bn2(self.conv2(x))) 155 | x = F.relu(self.bn3(self.conv3(x))) 156 | else: 157 | x = F.relu(self.conv1(x)) 158 | x = F.relu(self.conv2(x)) 159 | x = F.relu(self.conv3(x)) 160 | x = self.mp1(x) 161 | x = x.view(-1, 1024) 162 | 163 | if self.use_bn: 164 | x = F.relu(self.bn4(self.fc1(x))) 165 | x = F.relu(self.bn5(self.fc2(x))) 166 | else: 167 | x = F.relu(self.fc1(x)) 168 | x = F.relu(self.fc2(x)) 169 | x = self.fc3(x) 170 | 171 | iden = Variable(torch.from_numpy(np.eye(self.k).astype(np.float32))).view( 172 | 1, self.k*self.k).repeat(batchsize, 1) 173 | if x.is_cuda: 174 | iden = iden.cuda() 175 | x = x + iden 176 | x = x.view(-1, self.k, self.k) 177 | return x 178 | 179 | 180 | class PointNetfeat(nn.Module): 181 | def __init__(self, num_points=2500, global_feat=True, feature_transform=False, max_pool=True): 182 | super(PointNetfeat, self).__init__() 183 | self.stn = STN3d(num_points=num_points, k=3, use_bn=False) 184 | self.feature_trans = STN3d(num_points=num_points, k=64, use_bn=False) 185 | self.apply_feature_trans = feature_transform 186 | self.conv1 = torch.nn.Conv2d(1, 64, (1, 3)) 187 | self.conv2 = torch.nn.Conv2d(64, 64, (1, 1)) 188 | self.conv3 = torch.nn.Conv2d(64, 64, (1, 1)) 189 | self.conv4 = torch.nn.Conv2d(64, 128, (1, 1)) 190 | self.conv5 = torch.nn.Conv2d(128, 1024, (1, 1)) 191 | self.bn1 = nn.BatchNorm2d(64) 192 | self.bn2 = nn.BatchNorm2d(64) 193 | self.bn3 = nn.BatchNorm2d(64) 194 | self.bn4 = nn.BatchNorm2d(128) 195 | self.bn5 = nn.BatchNorm2d(1024) 196 | self.mp1 = torch.nn.MaxPool2d((num_points, 1), 1) 197 | self.num_points = num_points 198 | self.global_feat = global_feat 199 | self.max_pool = max_pool 200 | 201 | def forward(self, x): 202 | batchsize = x.size()[0] 203 | trans = self.stn(x) 204 | x = torch.matmul(torch.squeeze(x), trans) 205 | x = x.view(batchsize, 1, -1, 3) 206 | #x = x.transpose(2,1) 207 | #x = torch.bmm(x, trans) 208 | #x = x.transpose(2,1) 209 | x = F.relu(self.bn1(self.conv1(x))) 210 | x = F.relu(self.bn2(self.conv2(x))) 211 | pointfeat = x 212 | if self.apply_feature_trans: 213 | f_trans = self.feature_trans(x) 214 | x = torch.squeeze(x) 215 | if batchsize == 1: 216 | x = torch.unsqueeze(x, 0) 217 | x = torch.matmul(x.transpose(1, 2), f_trans) 218 | x = x.transpose(1, 2).contiguous() 219 | x = x.view(batchsize, 64, -1, 1) 220 | x = F.relu(self.bn3(self.conv3(x))) 221 | x = F.relu(self.bn4(self.conv4(x))) 222 | x = self.bn5(self.conv5(x)) 223 | if not self.max_pool: 224 | return x 225 | else: 226 | x = self.mp1(x) 227 | x = x.view(-1, 1024) 228 | if self.global_feat: 229 | return x, trans 230 | else: 231 | x = x.view(-1, 1024, 1).repeat(1, 1, self.num_points) 232 | return torch.cat([x, pointfeat], 1), trans 233 | 234 | 235 | class PointNetVlad(nn.Module): 236 | def __init__(self, num_points=2500, global_feat=True, feature_transform=False, max_pool=True, output_dim=1024): 237 | super(PointNetVlad, self).__init__() 238 | self.point_net = PointNetfeat(num_points=num_points, global_feat=global_feat, 239 | feature_transform=feature_transform, max_pool=max_pool) 240 | self.net_vlad = NetVLADLoupe(feature_size=1024, max_samples=num_points, cluster_size=64, 241 | output_dim=output_dim, gating=True, add_batch_norm=True, 242 | is_training=True) 243 | 244 | def forward(self, x): 245 | x = self.point_net(x) 246 | x = self.net_vlad(x) 247 | return x 248 | 249 | 250 | if __name__ == '__main__': 251 | num_points = 4096 252 | sim_data = Variable(torch.rand(44, 1, num_points, 3)) 253 | sim_data = sim_data.cuda() 254 | 255 | pnv = PointNetVlad.PointNetVlad(global_feat=True, feature_transform=True, max_pool=False, 256 | output_dim=256, num_points=num_points).cuda() 257 | pnv.train() 258 | out3 = pnv(sim_data) 259 | print('pnv', out3.size()) 260 | -------------------------------------------------------------------------------- /train_pointnetvlad.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import importlib 3 | import math 4 | import os 5 | import socket 6 | import sys 7 | 8 | import numpy as np 9 | from sklearn.neighbors import KDTree, NearestNeighbors 10 | 11 | import config as cfg 12 | import evaluate 13 | import loss.pointnetvlad_loss as PNV_loss 14 | import models.PointNetVlad as PNV 15 | import torch 16 | import torch.nn as nn 17 | from loading_pointclouds import * 18 | from tensorboardX import SummaryWriter 19 | from torch.autograd import Variable 20 | from torch.backends import cudnn 21 | 22 | os.environ["CUDA_VISIBLE_DEVICES"] = "3" 23 | 24 | 25 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 26 | sys.path.append(BASE_DIR) 27 | 28 | 29 | 30 | cudnn.enabled = True 31 | 32 | parser = argparse.ArgumentParser() 33 | parser.add_argument('--log_dir', default='log/', help='Log dir [default: log]') 34 | parser.add_argument('--results_dir', default='results/', 35 | help='results dir [default: results]') 36 | parser.add_argument('--positives_per_query', type=int, default=2, 37 | help='Number of potential positives in each training tuple [default: 2]') 38 | parser.add_argument('--negatives_per_query', type=int, default=18, 39 | help='Number of definite negatives in each training tuple [default: 18]') 40 | parser.add_argument('--max_epoch', type=int, default=20, 41 | help='Epoch to run [default: 20]') 42 | parser.add_argument('--batch_num_queries', type=int, default=2, 43 | help='Batch Size during training [default: 2]') 44 | parser.add_argument('--learning_rate', type=float, default=0.000005, 45 | help='Initial learning rate [default: 0.000005]') 46 | parser.add_argument('--momentum', type=float, default=0.9, 47 | help='Initial learning rate [default: 0.9]') 48 | parser.add_argument('--optimizer', default='adam', 49 | help='adam or momentum [default: adam]') 50 | parser.add_argument('--decay_step', type=int, default=200000, 51 | help='Decay step for lr decay [default: 200000]') 52 | parser.add_argument('--decay_rate', type=float, default=0.7, 53 | help='Decay rate for lr decay [default: 0.7]') 54 | parser.add_argument('--margin_1', type=float, default=0.5, 55 | help='Margin for hinge loss [default: 0.5]') 56 | parser.add_argument('--margin_2', type=float, default=0.2, 57 | help='Margin for hinge loss [default: 0.2]') 58 | parser.add_argument('--loss_function', default='quadruplet', choices=[ 59 | 'triplet', 'quadruplet'], help='triplet or quadruplet [default: quadruplet]') 60 | parser.add_argument('--loss_not_lazy', action='store_false', 61 | help='If present, do not use lazy variant of loss') 62 | parser.add_argument('--loss_ignore_zero_batch', action='store_true', 63 | help='If present, mean only batches with loss > 0.0') 64 | parser.add_argument('--triplet_use_best_positives', action='store_true', 65 | help='If present, use best positives, otherwise use hardest positives') 66 | parser.add_argument('--resume', action='store_true', 67 | help='If present, restore checkpoint and resume training') 68 | parser.add_argument('--dataset_folder', default='../../dataset/', 69 | help='PointNetVlad Dataset Folder') 70 | 71 | FLAGS = parser.parse_args() 72 | cfg.BATCH_NUM_QUERIES = FLAGS.batch_num_queries 73 | #cfg.EVAL_BATCH_SIZE = 12 74 | cfg.NUM_POINTS = 4096 75 | cfg.TRAIN_POSITIVES_PER_QUERY = FLAGS.positives_per_query 76 | cfg.TRAIN_NEGATIVES_PER_QUERY = FLAGS.negatives_per_query 77 | cfg.MAX_EPOCH = FLAGS.max_epoch 78 | cfg.BASE_LEARNING_RATE = FLAGS.learning_rate 79 | cfg.MOMENTUM = FLAGS.momentum 80 | cfg.OPTIMIZER = FLAGS.optimizer 81 | cfg.DECAY_STEP = FLAGS.decay_step 82 | cfg.DECAY_RATE = FLAGS.decay_rate 83 | cfg.MARGIN1 = FLAGS.margin_1 84 | cfg.MARGIN2 = FLAGS.margin_2 85 | cfg.FEATURE_OUTPUT_DIM = 256 86 | 87 | cfg.LOSS_FUNCTION = FLAGS.loss_function 88 | cfg.TRIPLET_USE_BEST_POSITIVES = FLAGS.triplet_use_best_positives 89 | cfg.LOSS_LAZY = FLAGS.loss_not_lazy 90 | cfg.LOSS_IGNORE_ZERO_BATCH = FLAGS.loss_ignore_zero_batch 91 | 92 | cfg.TRAIN_FILE = 'generating_queries/training_queries_baseline.pickle' 93 | cfg.TEST_FILE = 'generating_queries/test_queries_baseline.pickle' 94 | 95 | cfg.LOG_DIR = FLAGS.log_dir 96 | if not os.path.exists(cfg.LOG_DIR): 97 | os.mkdir(cfg.LOG_DIR) 98 | LOG_FOUT = open(os.path.join(cfg.LOG_DIR, 'log_train.txt'), 'w') 99 | LOG_FOUT.write(str(FLAGS) + '\n') 100 | 101 | cfg.RESULTS_FOLDER = FLAGS.results_dir 102 | 103 | cfg.DATASET_FOLDER = FLAGS.dataset_folder 104 | 105 | # Load dictionary of training queries 106 | TRAINING_QUERIES = get_queries_dict(cfg.TRAIN_FILE) 107 | TEST_QUERIES = get_queries_dict(cfg.TEST_FILE) 108 | 109 | cfg.BN_INIT_DECAY = 0.5 110 | cfg.BN_DECAY_DECAY_RATE = 0.5 111 | BN_DECAY_DECAY_STEP = float(cfg.DECAY_STEP) 112 | cfg.BN_DECAY_CLIP = 0.99 113 | 114 | HARD_NEGATIVES = {} 115 | TRAINING_LATENT_VECTORS = [] 116 | 117 | TOTAL_ITERATIONS = 0 118 | 119 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 120 | 121 | 122 | def get_bn_decay(batch): 123 | bn_momentum = cfg.BN_INIT_DECAY * \ 124 | (cfg.BN_DECAY_DECAY_RATE ** 125 | (batch * cfg.BATCH_NUM_QUERIES // BN_DECAY_DECAY_STEP)) 126 | return min(cfg.BN_DECAY_CLIP, 1 - bn_momentum) 127 | 128 | 129 | def log_string(out_str): 130 | LOG_FOUT.write(out_str + '\n') 131 | LOG_FOUT.flush() 132 | print(out_str) 133 | 134 | # learning rate halfed every 5 epoch 135 | 136 | 137 | def get_learning_rate(epoch): 138 | learning_rate = cfg.BASE_LEARNING_RATE * ((0.9) ** (epoch // 5)) 139 | learning_rate = max(learning_rate, 0.00001) # CLIP THE LEARNING RATE! 140 | return learning_rate 141 | 142 | 143 | def train(): 144 | global HARD_NEGATIVES, TOTAL_ITERATIONS 145 | bn_decay = get_bn_decay(0) 146 | #tf.summary.scalar('bn_decay', bn_decay) 147 | 148 | #loss = lazy_quadruplet_loss(q_vec, pos_vecs, neg_vecs, other_neg_vec, MARGIN1, MARGIN2) 149 | if cfg.LOSS_FUNCTION == 'quadruplet': 150 | loss_function = PNV_loss.quadruplet_loss 151 | else: 152 | loss_function = PNV_loss.triplet_loss_wrapper 153 | learning_rate = get_learning_rate(0) 154 | 155 | train_writer = SummaryWriter(os.path.join(cfg.LOG_DIR, 'train')) 156 | #test_writer = SummaryWriter(os.path.join(cfg.LOG_DIR, 'test')) 157 | 158 | model = PNV.PointNetVlad(global_feat=True, feature_transform=True, 159 | max_pool=False, output_dim=cfg.FEATURE_OUTPUT_DIM, num_points=cfg.NUM_POINTS) 160 | model = model.to(device) 161 | 162 | parameters = filter(lambda p: p.requires_grad, model.parameters()) 163 | 164 | if cfg.OPTIMIZER == 'momentum': 165 | optimizer = torch.optim.SGD( 166 | parameters, learning_rate, momentum=cfg.MOMENTUM) 167 | elif cfg.OPTIMIZER == 'adam': 168 | optimizer = torch.optim.Adam(parameters, learning_rate) 169 | else: 170 | optimizer = None 171 | exit(0) 172 | 173 | if FLAGS.resume: 174 | resume_filename = cfg.LOG_DIR + "checkpoint.pth.tar" 175 | print("Resuming From ", resume_filename) 176 | checkpoint = torch.load(resume_filename) 177 | saved_state_dict = checkpoint['state_dict'] 178 | starting_epoch = checkpoint['epoch'] 179 | TOTAL_ITERATIONS = starting_epoch * len(TRAINING_QUERIES) 180 | 181 | model.load_state_dict(saved_state_dict) 182 | optimizer.load_state_dict(checkpoint['optimizer']) 183 | else: 184 | starting_epoch = 0 185 | 186 | model = nn.DataParallel(model) 187 | 188 | LOG_FOUT.write(cfg.cfg_str()) 189 | LOG_FOUT.write("\n") 190 | LOG_FOUT.flush() 191 | 192 | for epoch in range(starting_epoch, cfg.MAX_EPOCH): 193 | print(epoch) 194 | print() 195 | log_string('**** EPOCH %03d ****' % (epoch)) 196 | sys.stdout.flush() 197 | 198 | train_one_epoch(model, optimizer, train_writer, loss_function, epoch) 199 | 200 | log_string('EVALUATING...') 201 | cfg.OUTPUT_FILE = cfg.RESULTS_FOLDER + 'results_' + str(epoch) + '.txt' 202 | eval_recall = evaluate.evaluate_model(model) 203 | log_string('EVAL RECALL: %s' % str(eval_recall)) 204 | 205 | train_writer.add_scalar("Val Recall", eval_recall, epoch) 206 | 207 | 208 | def train_one_epoch(model, optimizer, train_writer, loss_function, epoch): 209 | global HARD_NEGATIVES 210 | global TRAINING_LATENT_VECTORS, TOTAL_ITERATIONS 211 | 212 | is_training = True 213 | sampled_neg = 4000 214 | # number of hard negatives in the training tuple 215 | # which are taken from the sampled negatives 216 | num_to_take = 10 217 | 218 | # Shuffle train files 219 | train_file_idxs = np.arange(0, len(TRAINING_QUERIES.keys())) 220 | np.random.shuffle(train_file_idxs) 221 | 222 | for i in range(len(train_file_idxs)//cfg.BATCH_NUM_QUERIES): 223 | # for i in range (5): 224 | batch_keys = train_file_idxs[i * 225 | cfg.BATCH_NUM_QUERIES:(i+1)*cfg.BATCH_NUM_QUERIES] 226 | q_tuples = [] 227 | 228 | faulty_tuple = False 229 | no_other_neg = False 230 | for j in range(cfg.BATCH_NUM_QUERIES): 231 | if (len(TRAINING_QUERIES[batch_keys[j]]["positives"]) < cfg.TRAIN_POSITIVES_PER_QUERY): 232 | faulty_tuple = True 233 | break 234 | 235 | # no cached feature vectors 236 | if (len(TRAINING_LATENT_VECTORS) == 0): 237 | q_tuples.append( 238 | get_query_tuple(TRAINING_QUERIES[batch_keys[j]], cfg.TRAIN_POSITIVES_PER_QUERY, cfg.TRAIN_NEGATIVES_PER_QUERY, 239 | TRAINING_QUERIES, hard_neg=[], other_neg=True)) 240 | # q_tuples.append(get_rotated_tuple(TRAINING_QUERIES[batch_keys[j]],POSITIVES_PER_QUERY,NEGATIVES_PER_QUERY, TRAINING_QUERIES, hard_neg=[], other_neg=True)) 241 | # q_tuples.append(get_jittered_tuple(TRAINING_QUERIES[batch_keys[j]],POSITIVES_PER_QUERY,NEGATIVES_PER_QUERY, TRAINING_QUERIES, hard_neg=[], other_neg=True)) 242 | 243 | elif (len(HARD_NEGATIVES.keys()) == 0): 244 | query = get_feature_representation( 245 | TRAINING_QUERIES[batch_keys[j]]['query'], model) 246 | random.shuffle(TRAINING_QUERIES[batch_keys[j]]['negatives']) 247 | negatives = TRAINING_QUERIES[batch_keys[j] 248 | ]['negatives'][0:sampled_neg] 249 | hard_negs = get_random_hard_negatives( 250 | query, negatives, num_to_take) 251 | print(hard_negs) 252 | q_tuples.append( 253 | get_query_tuple(TRAINING_QUERIES[batch_keys[j]], cfg.TRAIN_POSITIVES_PER_QUERY, cfg.TRAIN_NEGATIVES_PER_QUERY, 254 | TRAINING_QUERIES, hard_negs, other_neg=True)) 255 | # q_tuples.append(get_rotated_tuple(TRAINING_QUERIES[batch_keys[j]],POSITIVES_PER_QUERY,NEGATIVES_PER_QUERY, TRAINING_QUERIES, hard_negs, other_neg=True)) 256 | # q_tuples.append(get_jittered_tuple(TRAINING_QUERIES[batch_keys[j]],POSITIVES_PER_QUERY,NEGATIVES_PER_QUERY, TRAINING_QUERIES, hard_negs, other_neg=True)) 257 | else: 258 | query = get_feature_representation( 259 | TRAINING_QUERIES[batch_keys[j]]['query'], model) 260 | random.shuffle(TRAINING_QUERIES[batch_keys[j]]['negatives']) 261 | negatives = TRAINING_QUERIES[batch_keys[j] 262 | ]['negatives'][0:sampled_neg] 263 | hard_negs = get_random_hard_negatives( 264 | query, negatives, num_to_take) 265 | hard_negs = list(set().union( 266 | HARD_NEGATIVES[batch_keys[j]], hard_negs)) 267 | print('hard', hard_negs) 268 | q_tuples.append( 269 | get_query_tuple(TRAINING_QUERIES[batch_keys[j]], cfg.TRAIN_POSITIVES_PER_QUERY, cfg.TRAIN_NEGATIVES_PER_QUERY, 270 | TRAINING_QUERIES, hard_negs, other_neg=True)) 271 | # q_tuples.append(get_rotated_tuple(TRAINING_QUERIES[batch_keys[j]],POSITIVES_PER_QUERY,NEGATIVES_PER_QUERY, TRAINING_QUERIES, hard_negs, other_neg=True)) 272 | # q_tuples.append(get_jittered_tuple(TRAINING_QUERIES[batch_keys[j]],POSITIVES_PER_QUERY,NEGATIVES_PER_QUERY, TRAINING_QUERIES, hard_negs, other_neg=True)) 273 | 274 | if (q_tuples[j][3].shape[0] != cfg.NUM_POINTS): 275 | no_other_neg = True 276 | break 277 | 278 | if(faulty_tuple): 279 | log_string('----' + str(i) + '-----') 280 | log_string('----' + 'FAULTY TUPLE' + '-----') 281 | continue 282 | 283 | if(no_other_neg): 284 | log_string('----' + str(i) + '-----') 285 | log_string('----' + 'NO OTHER NEG' + '-----') 286 | continue 287 | 288 | queries = [] 289 | positives = [] 290 | negatives = [] 291 | other_neg = [] 292 | for k in range(len(q_tuples)): 293 | queries.append(q_tuples[k][0]) 294 | positives.append(q_tuples[k][1]) 295 | negatives.append(q_tuples[k][2]) 296 | other_neg.append(q_tuples[k][3]) 297 | 298 | queries = np.array(queries, dtype=np.float32) 299 | queries = np.expand_dims(queries, axis=1) 300 | other_neg = np.array(other_neg, dtype=np.float32) 301 | other_neg = np.expand_dims(other_neg, axis=1) 302 | positives = np.array(positives, dtype=np.float32) 303 | negatives = np.array(negatives, dtype=np.float32) 304 | log_string('----' + str(i) + '-----') 305 | if (len(queries.shape) != 4): 306 | log_string('----' + 'FAULTY QUERY' + '-----') 307 | continue 308 | 309 | model.train() 310 | optimizer.zero_grad() 311 | 312 | output_queries, output_positives, output_negatives, output_other_neg = run_model( 313 | model, queries, positives, negatives, other_neg) 314 | loss = loss_function(output_queries, output_positives, output_negatives, output_other_neg, cfg.MARGIN_1, cfg.MARGIN_2, use_min=cfg.TRIPLET_USE_BEST_POSITIVES, lazy=cfg.LOSS_LAZY, ignore_zero_loss=cfg.LOSS_IGNORE_ZERO_BATCH) 315 | loss.backward() 316 | optimizer.step() 317 | 318 | log_string('batch loss: %f' % loss) 319 | train_writer.add_scalar("Loss", loss.cpu().item(), TOTAL_ITERATIONS) 320 | TOTAL_ITERATIONS += cfg.BATCH_NUM_QUERIES 321 | 322 | # EVALLLL 323 | 324 | if (epoch > 5 and i % (1400 // cfg.BATCH_NUM_QUERIES) == 29): 325 | TRAINING_LATENT_VECTORS = get_latent_vectors( 326 | model, TRAINING_QUERIES) 327 | print("Updated cached feature vectors") 328 | 329 | if (i % (6000 // cfg.BATCH_NUM_QUERIES) == 101): 330 | if isinstance(model, nn.DataParallel): 331 | model_to_save = model.module 332 | else: 333 | model_to_save = model 334 | save_name = cfg.LOG_DIR + cfg.MODEL_FILENAME 335 | torch.save({ 336 | 'epoch': epoch, 337 | 'iter': TOTAL_ITERATIONS, 338 | 'state_dict': model_to_save.state_dict(), 339 | 'optimizer': optimizer.state_dict(), 340 | }, 341 | save_name) 342 | print("Model Saved As " + save_name) 343 | 344 | 345 | def get_feature_representation(filename, model): 346 | model.eval() 347 | queries = load_pc_files([filename]) 348 | queries = np.expand_dims(queries, axis=1) 349 | # if(BATCH_NUM_QUERIES-1>0): 350 | # fake_queries=np.zeros((BATCH_NUM_QUERIES-1,1,NUM_POINTS,3)) 351 | # q=np.vstack((queries,fake_queries)) 352 | # else: 353 | # q=queries 354 | with torch.no_grad(): 355 | q = torch.from_numpy(queries).float() 356 | q = q.to(device) 357 | output = model(q) 358 | output = output.detach().cpu().numpy() 359 | output = np.squeeze(output) 360 | model.train() 361 | return output 362 | 363 | 364 | def get_random_hard_negatives(query_vec, random_negs, num_to_take): 365 | global TRAINING_LATENT_VECTORS 366 | 367 | latent_vecs = [] 368 | for j in range(len(random_negs)): 369 | latent_vecs.append(TRAINING_LATENT_VECTORS[random_negs[j]]) 370 | 371 | latent_vecs = np.array(latent_vecs) 372 | nbrs = KDTree(latent_vecs) 373 | distances, indices = nbrs.query(np.array([query_vec]), k=num_to_take) 374 | hard_negs = np.squeeze(np.array(random_negs)[indices[0]]) 375 | hard_negs = hard_negs.tolist() 376 | return hard_negs 377 | 378 | 379 | def get_latent_vectors(model, dict_to_process): 380 | train_file_idxs = np.arange(0, len(dict_to_process.keys())) 381 | 382 | batch_num = cfg.BATCH_NUM_QUERIES * \ 383 | (1 + cfg.TRAIN_POSITIVES_PER_QUERY + cfg.TRAIN_NEGATIVES_PER_QUERY + 1) 384 | q_output = [] 385 | 386 | model.eval() 387 | 388 | for q_index in range(len(train_file_idxs)//batch_num): 389 | file_indices = train_file_idxs[q_index * 390 | batch_num:(q_index+1)*(batch_num)] 391 | file_names = [] 392 | for index in file_indices: 393 | file_names.append(dict_to_process[index]["query"]) 394 | queries = load_pc_files(file_names) 395 | 396 | feed_tensor = torch.from_numpy(queries).float() 397 | feed_tensor = feed_tensor.unsqueeze(1) 398 | feed_tensor = feed_tensor.to(device) 399 | with torch.no_grad(): 400 | out = model(feed_tensor) 401 | 402 | out = out.detach().cpu().numpy() 403 | out = np.squeeze(out) 404 | 405 | q_output.append(out) 406 | 407 | q_output = np.array(q_output) 408 | if(len(q_output) != 0): 409 | q_output = q_output.reshape(-1, q_output.shape[-1]) 410 | 411 | # handle edge case 412 | for q_index in range((len(train_file_idxs) // batch_num * batch_num), len(dict_to_process.keys())): 413 | index = train_file_idxs[q_index] 414 | queries = load_pc_files([dict_to_process[index]["query"]]) 415 | queries = np.expand_dims(queries, axis=1) 416 | 417 | # if (BATCH_NUM_QUERIES - 1 > 0): 418 | # fake_queries = np.zeros((BATCH_NUM_QUERIES - 1, 1, NUM_POINTS, 3)) 419 | # q = np.vstack((queries, fake_queries)) 420 | # else: 421 | # q = queries 422 | 423 | #fake_pos = np.zeros((BATCH_NUM_QUERIES, POSITIVES_PER_QUERY, NUM_POINTS, 3)) 424 | #fake_neg = np.zeros((BATCH_NUM_QUERIES, NEGATIVES_PER_QUERY, NUM_POINTS, 3)) 425 | #fake_other_neg = np.zeros((BATCH_NUM_QUERIES, 1, NUM_POINTS, 3)) 426 | #o1, o2, o3, o4 = run_model(model, q, fake_pos, fake_neg, fake_other_neg) 427 | with torch.no_grad(): 428 | queries_tensor = torch.from_numpy(queries).float() 429 | o1 = model(queries_tensor) 430 | 431 | output = o1.detach().cpu().numpy() 432 | output = np.squeeze(output) 433 | if (q_output.shape[0] != 0): 434 | q_output = np.vstack((q_output, output)) 435 | else: 436 | q_output = output 437 | 438 | model.train() 439 | print(q_output.shape) 440 | return q_output 441 | 442 | 443 | def run_model(model, queries, positives, negatives, other_neg, require_grad=True): 444 | queries_tensor = torch.from_numpy(queries).float() 445 | positives_tensor = torch.from_numpy(positives).float() 446 | negatives_tensor = torch.from_numpy(negatives).float() 447 | other_neg_tensor = torch.from_numpy(other_neg).float() 448 | feed_tensor = torch.cat( 449 | (queries_tensor, positives_tensor, negatives_tensor, other_neg_tensor), 1) 450 | feed_tensor = feed_tensor.view((-1, 1, cfg.NUM_POINTS, 3)) 451 | feed_tensor.requires_grad_(require_grad) 452 | feed_tensor = feed_tensor.to(device) 453 | if require_grad: 454 | output = model(feed_tensor) 455 | else: 456 | with torch.no_grad(): 457 | output = model(feed_tensor) 458 | output = output.view(cfg.BATCH_NUM_QUERIES, -1, cfg.FEATURE_OUTPUT_DIM) 459 | o1, o2, o3, o4 = torch.split( 460 | output, [1, cfg.TRAIN_POSITIVES_PER_QUERY, cfg.TRAIN_NEGATIVES_PER_QUERY, 1], dim=1) 461 | 462 | return o1, o2, o3, o4 463 | 464 | 465 | if __name__ == "__main__": 466 | train() 467 | --------------------------------------------------------------------------------