├── README.md ├── datasets ├── __pycache__ │ ├── cla_scalable.cpython-38.pyc │ ├── data_utils.cpython-36.pyc │ ├── data_utils.cpython-38.pyc │ ├── scalable_dataset.cpython-38.pyc │ ├── scannetpr_dataset.cpython-38.pyc │ ├── test_dataset.cpython-38.pyc │ ├── train_classification_datasets.cpython-38.pyc │ ├── train_datasets.cpython-36.pyc │ └── train_datasets.cpython-38.pyc ├── data_utils.py ├── test_dataset.py ├── train_datasets.py └── triplets │ ├── __pycache__ │ ├── base_dataset.cpython-38.pyc │ ├── dataloader.cpython-38.pyc │ ├── mulran_dataset.cpython-38.pyc │ ├── pointcloud_dataset.cpython-38.pyc │ └── samplers.cpython-38.pyc │ ├── base_dataset.py │ ├── dataloader.py │ ├── mulran_dataset.py │ ├── pointcloud_dataset.py │ ├── samplers.py │ └── scannetpr_dataset.py ├── eval.sh ├── eval ├── __pycache__ │ ├── eval_config.cpython-38.pyc │ ├── eval_sequence.cpython-38.pyc │ └── plot_loop.cpython-38.pyc ├── eval_config.py ├── eval_sequence.py ├── indoor_localization.py ├── loop_closure.py └── relocalization.py ├── libs ├── __init__.py ├── __pycache__ │ └── __init__.cpython-36.pyc └── pointops │ ├── __init__.py │ ├── functions │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ └── pointops.cpython-36.pyc │ └── pointops.py │ ├── setup.py │ └── src │ ├── __init__.py │ ├── ballquery │ ├── ballquery_cuda.cpp │ ├── ballquery_cuda_kernel.cu │ └── ballquery_cuda_kernel.h │ ├── cuda_utils.h │ ├── featuredistribute │ ├── featuredistribute_cuda.cpp │ ├── featuredistribute_cuda_kernel.cu │ └── featuredistribute_cuda_kernel.h │ ├── grouping │ ├── grouping_cuda.cpp │ ├── grouping_cuda_kernel.cu │ └── grouping_cuda_kernel.h │ ├── grouping_int │ ├── grouping_int_cuda.cpp │ ├── grouping_int_cuda_kernel.cu │ └── grouping_int_cuda_kernel.h │ ├── interpolation │ ├── interpolation_cuda.cpp │ ├── interpolation_cuda_kernel.cu │ └── interpolation_cuda_kernel.h │ ├── knnquery │ ├── __init__.py │ ├── knnquery_cuda.cpp │ ├── knnquery_cuda_kernel.cu │ └── knnquery_cuda_kernel.h │ ├── labelstat │ ├── labelstat_cuda.cpp │ ├── labelstat_cuda_kernel.cu │ └── labelstat_cuda_kernel.h │ ├── pointops_api.cpp │ └── sampling │ ├── sampling_cuda.cpp │ ├── sampling_cuda_kernel.cu │ └── sampling_cuda_kernel.h ├── loss ├── __pycache__ │ ├── cosface_loss.cpython-36.pyc │ ├── cosface_loss.cpython-38.pyc │ ├── incremental_loss.cpython-38.pyc │ ├── triplet_loss.cpython-36.pyc │ └── triplet_loss.cpython-38.pyc ├── cosface_loss.py ├── incremental_loss.py └── triplet_loss.py ├── models ├── MinkLoc3dv2 │ ├── __pycache__ │ │ ├── mink_params.cpython-38.pyc │ │ ├── minkfpn.cpython-38.pyc │ │ ├── minkloc.cpython-38.pyc │ │ └── resnet.cpython-38.pyc │ ├── layers │ │ ├── __pycache__ │ │ │ ├── eca_block.cpython-36.pyc │ │ │ ├── eca_block.cpython-38.pyc │ │ │ ├── netvlad.cpython-38.pyc │ │ │ ├── pooling.cpython-38.pyc │ │ │ └── pooling_wrapper.cpython-38.pyc │ │ ├── eca_block.py │ │ ├── netvlad.py │ │ ├── pooling.py │ │ └── pooling_wrapper.py │ ├── losses │ │ ├── __pycache__ │ │ │ ├── loss.cpython-38.pyc │ │ │ ├── loss_utils.cpython-38.pyc │ │ │ └── truncated_smoothap.cpython-38.pyc │ │ ├── loss.py │ │ ├── loss_utils.py │ │ └── truncated_smoothap.py │ ├── mink_params.py │ ├── minkfpn.py │ ├── minkloc.py │ ├── minkloc3dv2.txt │ └── resnet.py ├── PointNetVlad.py ├── __pycache__ │ ├── PointNetVlad.cpython-38.pyc │ ├── loupe.cpython-38.pyc │ ├── model_factory.cpython-38.pyc │ ├── pptnet_v2.cpython-38.pyc │ └── pt_util.cpython-38.pyc ├── loupe.py ├── model_factory.py ├── pointnet_util.py ├── pptnet_v2.py └── pt_util.py ├── train.py ├── train.sh └── util ├── __pycache__ ├── computer_overlap.cpython-38.pyc ├── loading_pointclouds.cpython-36.pyc ├── loading_pointclouds.cpython-38.pyc ├── parser.cpython-36.pyc ├── parser.cpython-38.pyc ├── pointnetvlad_loss.cpython-36.pyc ├── pointnetvlad_loss.cpython-38.pyc ├── pt_util.cpython-36.pyc ├── pt_util.cpython-38.pyc ├── scannet_config.cpython-38.pyc ├── util.cpython-36.pyc └── util.cpython-38.pyc ├── loading_pointclouds.py ├── parser.py ├── pointnetvlad_loss.py ├── pt_util.py ├── scannet_config.py └── util.py /README.md: -------------------------------------------------------------------------------- 1 | # Look At the Whole Scene: General Point Cloud Place Recognition by Classification Proxy 2 | 3 | ## Abtract 4 | 5 | Deep learning models centered on retrieval have made significant strides in point cloud place recognition. However, existing approaches struggle to generate discriminative global descriptors and often rely on labor-intensive negative sample mining. Such constraints limit their usability in dynamic and open-world scenarios. To address these challenges, we introduce LAWS, a pioneering classification-centric neural framework that emphasizes looking at the whole scene for superior point cloud descriptor extraction. Central to our approach is the space partitioning design, constructed to provide holistic scene supervision, ensuring the comprehensive learning of scene features. To counteract potential ambiguities arising from the single orthogonal partition boundary, a complementary mechanism of repartitioning space diagonally is specifically designed to dispel classification uncertainties. Under the enhanced partitioning mechanism, the space is separated into several classes and groups. Furthermore, to prevent knowledge forgetting between different groups, a special training strategy is employed, allowing for distinct training of each group. The extensive experiments, encompassing both indoor and outdoor settings and different tasks, validate the generality of LAWS. It not only outperforms contemporary methods but also demonstrates a profound generalization ability across various unseen environments and sensor modalities. Our method achieves a 2.6% higher average top-1 recall on Oxford RobotCar Dataset and a 7.8% higher average recall when generalized to In-house Dataset compared with retrieval-based methods. Furthermore, LAWS also outperforms retrieval-based methods in terms of F1 score, with improvements of 12.7 and 29.2 on the MulRan and KITTI datasets, respectively. Notably, the average localization accuracy of LAWS in indoor environments reached about 68.1%. Moreover, the scalability and efficiency places LAWS in a leading position for continuous exploration and long-term autonomy. 6 | 7 | 8 | 9 | ## Datasets 10 | 11 | - Oxford RobotCar Dataset 12 | - In-house Datasets 13 | - university sector (U.S.) 14 | - residential area (R.A.) 15 | - business district (B.D.) 16 | - MulRan 17 | - KITTI odometry 18 | - ScannetPR 19 | 20 | Following [PointNetVLAD](https://arxiv.org/abs/1804.03492) the Oxford and In-house datasets can be [downloaded](https://drive.google.com/open?id=1H9Ep76l8KkUpwILY-13owsEMbVCYTmyx) here. MulRan dataset 21 | 22 | 23 | ## Environment and Dependencies 24 | 25 | 26 | This project has been tested using Python 3.8.19 with PyTorch 1.10.2 and MinkowskiEngine 0.5.4 on Ubuntu 20.04 with CUDA 11.4. Main dependencies include: 27 | 28 | The following packages are required: 29 | 30 | - PyTorch (version 1.10.1) 31 | - Torchvision (version 0.11.2) 32 | - MinkowskiEngine (version 0.5.4) 33 | - pandas 34 | - tqdm 35 | - scipy 36 | - tensorboardX 37 | 38 | Set up the requirments as follows: 39 | 40 | 1. Create conda environment with python: 41 | ``` 42 | conda create -n LAWS python=3.8 43 | conda activate LAWS 44 | ``` 45 | 2. Install PyTorch with suitable cudatoolkit version. 46 | 47 | 3. Install [MinkowskiEngine](https://github.com/NVIDIA/MinkowskiEngine) for MinkLoc3dv2 48 | ``` 49 | For example, if you want local MinkowskiEngine 50 | export CUDA_HOME=/usr/local/cuda-11.X 51 | git clone https://github.com/NVIDIA/MinkowskiEngine.git 52 | cd MinkowskiEngine 53 | python setup.py install --blas_include_dirs=${CONDA_PREFIX}/include --blas=openblas 54 | ``` 55 | 4. Build the pointops for PPTNet 56 | ``` 57 | cd libs/pointops && python setup.py install && cd ../../ 58 | ``` 59 | 60 | ## Training & Evaluating 61 | 1. Train the network 62 | ``` 63 | sh train.sh 64 | ``` 65 | 2. Evaluate the network 66 | ``` 67 | sh eval.sh 68 | ``` 69 | 70 | 71 | ## Acknowledgement 72 | 73 | Our code refers to [PointNetVLAD](https://github.com/mikacuy/pointnetvlad), [MinkLoc3Dv2](https://github.com/jac99/MinkLoc3Dv2), [PPT-Net](https://github.com/fpthink/PPT-Net) and [CosPlace](https://github.com/gmberton/CosPlace). 74 | 75 | 76 | -------------------------------------------------------------------------------- /datasets/__pycache__/cla_scalable.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/cla_scalable.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/data_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/data_utils.cpython-36.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/data_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/data_utils.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/scalable_dataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/scalable_dataset.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/scannetpr_dataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/scannetpr_dataset.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/test_dataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/test_dataset.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/train_classification_datasets.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/train_classification_datasets.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/train_datasets.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/train_datasets.cpython-36.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/train_datasets.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/__pycache__/train_datasets.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/data_utils.py: -------------------------------------------------------------------------------- 1 | 2 | def get_sets_dict(filename): 3 | #[key_dataset:{key_pointcloud:{'query':file,'northing':value,'easting':value}},key_dataset:{key_pointcloud:{'query':file,'northing':value,'easting':value}}, ...} 4 | with open(filename, 'rb') as handle: 5 | trajectories = pickle.load(handle) 6 | return trajectories 7 | -------------------------------------------------------------------------------- /datasets/triplets/__pycache__/base_dataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/triplets/__pycache__/base_dataset.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/triplets/__pycache__/dataloader.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/triplets/__pycache__/dataloader.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/triplets/__pycache__/mulran_dataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/triplets/__pycache__/mulran_dataset.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/triplets/__pycache__/pointcloud_dataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/triplets/__pycache__/pointcloud_dataset.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/triplets/__pycache__/samplers.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/datasets/triplets/__pycache__/samplers.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/triplets/base_dataset.py: -------------------------------------------------------------------------------- 1 | # Base dataset classes, inherited by dataset-specific classes 2 | import os 3 | import pickle 4 | from typing import List 5 | from typing import Dict 6 | import torch 7 | import numpy as np 8 | from torch.utils.data import Dataset 9 | import logging, json 10 | import glob 11 | 12 | class TrainingTuple: 13 | # Tuple describing an element for training/validation 14 | def __init__(self, id: int, timestamp: int, rel_scan_filepath: str, positives: np.ndarray, 15 | non_negatives: np.ndarray, position: np.ndarray): 16 | # id: element id (ids start from 0 and are consecutive numbers) 17 | # ts: timestamp 18 | # rel_scan_filepath: relative path to the scan 19 | # positives: sorted ndarray of positive elements id 20 | # negatives: sorted ndarray of elements id 21 | # position: x, y position in meters (northing, easting) 22 | assert position.shape == (2,) 23 | 24 | self.id = id 25 | self.timestamp = timestamp 26 | self.rel_scan_filepath = rel_scan_filepath 27 | self.positives = positives 28 | self.non_negatives = non_negatives 29 | self.position = position 30 | 31 | class TrainingDataset(Dataset): 32 | def __init__(self, dataset_path, transform=None, set_transform=None): 33 | 34 | self.transform = transform 35 | self.set_transform = set_transform 36 | self.files = [] 37 | self.root = dataset_path 38 | logging.info("Initializing MulRanTupleDataset") 39 | logging.info(f"Loading the data from {self.root}") 40 | 41 | sequences = ['DCC/DCC_01', 'DCC/DCC_02', 42 | 'Riverside/Riverside_01', 'Riverside/Riverside_03'] 43 | tuple_dir = os.path.join(os.path.dirname( 44 | __file__), '../../configs/mulran_tuples/') 45 | # tuple_dir = '/configs/mulran_tuples/' 46 | self.dict_3m = json.load(open(tuple_dir + 'positive_sequence_D-3_T-0.json', "r")) 47 | self.dict_20m = json.load( 48 | open(tuple_dir + 'positive_sequence_D-20_T-0.json', "r")) 49 | self.mulran_seq_lens = {"DCC/DCC_01": 5542, "DCC/DCC_02": 7561, "DCC/DCC_03": 7479, 50 | "KAIST/KAIST_01": 8226, "KAIST/KAIST_02": 8941, "KAIST/KAIST_03": 8629, 51 | "Sejong/Sejong_01": 28779, "Sejong/Sejong_02": 27494, "Sejong/Sejong_03": 27215, 52 | "Riverside/Riverside_01": 5537, "Riverside/Riverside_02": 8157, "Riverside/Riverside_03": 10476} 53 | 54 | for drive_id in sequences: 55 | sequence_path = self.root + drive_id + '/Downsample/' 56 | fnames = sorted(glob.glob(os.path.join(sequence_path, '*.bin'))) 57 | assert len( 58 | fnames) > 0, f"Make sure that the path {self.root} has data {drive_id}" 59 | inames = sorted([int(os.path.split(fname)[-1][:-4]) 60 | for fname in fnames]) 61 | 62 | for query_id, start_time in enumerate(inames): 63 | positives = self.get_positives(drive_id, query_id) 64 | non_negatives = self.get_non_negatives(drive_id, query_id) 65 | self.files.append((drive_id, query_id, positives, non_negatives)) 66 | 67 | print('{} queries in the dataset'.format(len(self.files))) 68 | 69 | def __len__(self): 70 | return len(self.files) 71 | 72 | def get_velodyne_fn(self, drive_id, query_id): 73 | sequence_path = self.root + drive_id + '/Downsample/' 74 | fname = sorted(glob.glob(os.path.join( 75 | sequence_path, '*.bin')))[query_id] 76 | return fname 77 | 78 | def get_pointcloud_tensor(self, drive_id, pc_id): 79 | fname = self.get_velodyne_fn(drive_id, pc_id) 80 | pc = np.fromfile(fname, dtype=np.float64).reshape(-1, 3) 81 | 82 | if(pc.shape[0] != 4096) or (pc.shape[1] != 3): 83 | print("Error in pointcloud shape", pc.shape) 84 | return np.array([]) 85 | return pc 86 | 87 | def __getitem__(self, idx): 88 | # Load point cloud and apply transform 89 | drive_id, query_id = self.files[idx][0], self.files[idx][1] 90 | # positive_ids, non_negatives_ids = self.files[idx][2], self.files[idx][3] 91 | 92 | query = self.get_pointcloud_tensor(drive_id, query_id) 93 | 94 | return query, idx 95 | 96 | def get_positives(self, sq, index): 97 | assert sq in self.dict_3m.keys(), f"Error: Sequence {sq} not in json." 98 | sq_1 = self.dict_3m[sq] 99 | if str(int(index)) in sq_1: 100 | positives = sq_1[str(int(index))] 101 | else: 102 | positives = [] 103 | return positives 104 | 105 | def get_negatives(self, sq, index): 106 | assert sq in self.dict_20m.keys(), f"Error: Sequence {sq} not in json." 107 | sq_2 = self.dict_20m[sq] 108 | all_ids = set(np.arange(self.mulran_seq_lens[sq])) 109 | neg_set_inv = sq_2[str(int(index))] 110 | neg_set = all_ids.difference(neg_set_inv) 111 | negatives = list(neg_set) 112 | if index in negatives: 113 | negatives.remove(index) 114 | return negatives 115 | 116 | def get_non_negatives(self, sq, index): 117 | assert sq in self.dict_20m.keys(), f"Error: Sequence {sq} not in json." 118 | sq_3 = self.dict_20m[sq] 119 | if str(int(index)) in sq_3: 120 | non_negatives = sq_3[str(int(index))] 121 | else: 122 | non_negatives = [] 123 | return non_negatives 124 | 125 | 126 | class PointCloudLoader: 127 | def __init__(self): 128 | # remove_zero_points: remove points with all zero coordinates 129 | # remove_ground_plane: remove points on ground plane level and below 130 | # ground_plane_level: ground plane level 131 | self.remove_zero_points = True 132 | self.remove_ground_plane = True 133 | self.ground_plane_level = None 134 | self.set_properties() 135 | 136 | def set_properties(self): 137 | # Set point cloud properties, such as ground_plane_level. Must be defined in inherited classes. 138 | raise NotImplementedError('set_properties must be defined in inherited classes') 139 | 140 | def __call__(self, file_pathname): 141 | # Reads the point cloud from a disk and preprocess (optional removal of zero points and points on the ground 142 | # plane and below 143 | # file_pathname: relative file path 144 | assert os.path.exists(file_pathname), f"Cannot open point cloud: {file_pathname}" 145 | pc = self.read_pc(file_pathname) 146 | assert pc.shape[1] == 3 147 | 148 | if self.remove_zero_points: 149 | mask = np.all(np.isclose(pc, 0), axis=1) 150 | pc = pc[~mask] 151 | 152 | if self.remove_ground_plane: 153 | mask = pc[:, 2] > self.ground_plane_level 154 | pc = pc[mask] 155 | 156 | return pc 157 | 158 | def read_pc(self, file_pathname: str) -> np.ndarray: 159 | # Reads the point cloud without pre-processing 160 | raise NotImplementedError("read_pc must be overloaded in an inheriting class") 161 | 162 | class PNVTrainingDataset(TrainingDataset): 163 | def __init__(self, *args, **kwargs): 164 | super().__init__(*args, **kwargs) 165 | self.pc_loader = PNVPointCloudLoader() 166 | -------------------------------------------------------------------------------- /datasets/triplets/dataloader.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import logging 4 | from datasets.triplets.mulran_dataset import MulRanTupleDataset 5 | from models.MinkLoc3dv2.mink_params import TrainingParams 6 | from datasets.triplets.base_dataset import TrainingDataset 7 | from datasets.triplets.samplers import BatchSampler 8 | from torch.utils.data.sampler import Sampler 9 | import numpy as np 10 | import MinkowskiEngine as ME 11 | from torch.utils.data import DataLoader 12 | from datasets.triplets.scannetpr_dataset import ScannetTripleDataset 13 | 14 | def in_sorted_array(e: int, array: np.ndarray) -> bool: 15 | pos = np.searchsorted(array, e) 16 | if pos == len(array) or pos == -1: 17 | return False 18 | else: 19 | return array[pos] == e 20 | 21 | class RandomSampler(Sampler): 22 | """Samples elements randomly, without replacement. 23 | Arguments: 24 | data_source (Dataset): dataset to sample from 25 | shuffle: use random permutation 26 | """ 27 | 28 | def __init__(self, data_source, shuffle=False): 29 | self.data_source = data_source 30 | self.shuffle = shuffle 31 | self.reset_permutation() 32 | 33 | def reset_permutation(self): 34 | perm = len(self.data_source) 35 | if self.shuffle: 36 | perm = torch.randperm(perm) 37 | else: 38 | perm = torch.arange(perm) 39 | self._perm = perm.tolist() 40 | 41 | def __iter__(self): 42 | return self 43 | 44 | def __next__(self): 45 | if len(self._perm) == 0: 46 | self.reset_permutation() 47 | return self._perm.pop(0) 48 | 49 | def __len__(self): 50 | return len(self.data_source) 51 | 52 | def collate_tuple(list_data): 53 | queries = [] 54 | positives = [] 55 | negatives = [] 56 | other_neg = [] 57 | for k in range(len(list_data)): 58 | queries.append(list_data[k][0]) 59 | positives.append(list_data[k][1]) 60 | negatives.append(list_data[k][2]) 61 | other_neg.append(list_data[k][3]) 62 | 63 | queries, positives, negatives, other_neg = np.array(queries), np.array(positives), np.array(negatives), np.array(other_neg) 64 | queries_tensor = torch.from_numpy(queries).unsqueeze(1).float() 65 | positives_tensor = torch.from_numpy(positives).float() 66 | negatives_tensor = torch.from_numpy(negatives).float() 67 | other_neg_tensor = torch.from_numpy(other_neg).unsqueeze(1).float() 68 | feed_tensor = torch.cat( 69 | (queries_tensor, positives_tensor, negatives_tensor, other_neg_tensor), 1) 70 | feed_tensor = feed_tensor.view((-1, 1, 4096, 3)).squeeze(1) 71 | # for turple_data in inputs: 72 | # for data in turple_data: 73 | # if isinstance(data, np.ndarray): 74 | # outputs.append(data) 75 | # elif isinstance(data, list): 76 | # outputs.extend(data) 77 | 78 | # outputs = np.array(outputs) 79 | 80 | return feed_tensor 81 | 82 | 83 | def make_datasets(params: TrainingParams, validation: bool = True): 84 | # Create training and validation datasets 85 | datasets = {} 86 | train_set_transform = None 87 | 88 | # PoinNetVLAD datasets (RobotCar and Inhouse) 89 | # PNV datasets have their own transform 90 | train_transform = None 91 | datasets['train'] = TrainingDataset(params.dataset_folder, transform=train_transform, set_transform=train_set_transform) 92 | if validation: 93 | datasets['val'] = TrainingDataset(params.dataset_folder, params.val_file) 94 | 95 | return datasets 96 | 97 | def make_dataloaders(params: TrainingParams, validation=False): 98 | 99 | ###################################################################### 100 | datasets = make_datasets(params, validation=validation) 101 | 102 | dataloders = {} 103 | train_sampler = BatchSampler(datasets['train'], batch_size=params.batch_size, 104 | batch_size_limit=params.batch_size_limit, 105 | batch_expansion_rate=params.batch_expansion_rate) 106 | 107 | # Collate function collates items into a batch and applies a 'set transform' on the entire batch 108 | quantizer = params.model_params.quantizer 109 | train_collate_fn = make_collate_fn(datasets['train'], quantizer, params.batch_split_size) 110 | dataloders['train'] = DataLoader(datasets['train'], batch_sampler=train_sampler, 111 | collate_fn=train_collate_fn, num_workers=params.num_workers, 112 | pin_memory=True) 113 | if validation and 'val' in datasets: 114 | val_collate_fn = make_collate_fn(datasets['val'], quantizer, params.batch_split_size) 115 | val_sampler = BatchSampler(datasets['val'], batch_size=params.val_batch_size) 116 | # Collate function collates items into a batch and applies a 'set transform' on the entire batch 117 | # Currently validation dataset has empty set_transform function, but it may change in the future 118 | dataloders['val'] = DataLoader(datasets['val'], batch_sampler=val_sampler, collate_fn=val_collate_fn, 119 | num_workers=params.num_workers, pin_memory=True) 120 | 121 | return dataloders 122 | 123 | 124 | def make_collate_fn(dataset: TrainingDataset, quantizer, batch_split_size=None): 125 | # quantizer: converts to polar (when polar coords are used) and quantizes 126 | # batch_split_size: if not None, splits the batch into a list of multiple mini-batches with batch_split_size elems 127 | def collate_fn(data_list): 128 | # Constructs a batch object 129 | clouds = [e[0] for e in data_list] 130 | labels = [e[1] for e in data_list] 131 | 132 | if dataset.set_transform is not None: 133 | # Apply the same transformation on all dataset elements 134 | lens = [len(cloud) for cloud in clouds] 135 | clouds = torch.cat(clouds, dim=0) 136 | clouds = dataset.set_transform(clouds) 137 | clouds = clouds.split(lens) 138 | 139 | # Compute positives and negatives mask 140 | # dataset.queries[label]['positives'] is bitarray 141 | positives_mask = [[in_sorted_array(e, dataset.files[label][2]) for e in labels] for label in labels] 142 | negatives_mask = [[not in_sorted_array(e, dataset.files[label][3]) for e in labels] for label in labels] 143 | positives_mask = torch.tensor(positives_mask) 144 | negatives_mask = torch.tensor(negatives_mask) 145 | 146 | # Convert to polar (when polar coords are used) and quantize 147 | # Use the first value returned by quantizer 148 | coords = [quantizer(e)[0] for e in clouds] 149 | 150 | if batch_split_size is None or batch_split_size == 0: 151 | coords = ME.utils.batched_coordinates(coords) 152 | # Assign a dummy feature equal to 1 to each point 153 | feats = torch.ones((coords.shape[0], 1), dtype=torch.float32) 154 | batch = {'coords': coords, 'features': feats} 155 | 156 | else: 157 | # Split the batch into chunks 158 | batch = [] 159 | for i in range(0, len(coords), batch_split_size): 160 | temp = coords[i:i + batch_split_size] 161 | c = ME.utils.batched_coordinates(temp) 162 | f = torch.ones((c.shape[0], 1), dtype=torch.float32) 163 | minibatch = {'coords': c, 'features': f} 164 | batch.append(minibatch) 165 | 166 | # Returns (batch_size, n_points, 3) tensor and positives_mask and negatives_mask which are 167 | # batch_size x batch_size boolean tensors 168 | #return batch, positives_mask, negatives_mask, torch.tensor(sampled_positive_ndx), torch.tensor(relative_poses) 169 | return batch, positives_mask, negatives_mask 170 | 171 | return collate_fn 172 | 173 | 174 | 175 | def make_data_loader_for_scannet(args): 176 | dset = ScannetTripleDataset(args,'training') 177 | sampler = RandomSampler(dset, shuffle=True) 178 | data_loader = torch.utils.data.DataLoader(dset, 179 | num_workers=16, 180 | batch_size=args["BATCH_NUM_QUERIES"], 181 | sampler=sampler, 182 | pin_memory=True) 183 | return data_loader 184 | -------------------------------------------------------------------------------- /datasets/triplets/mulran_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import glob 4 | import random 5 | import numpy as np 6 | import logging 7 | import json 8 | import csv 9 | import torch 10 | 11 | sys.path.append(os.path.join(os.path.dirname(__file__), '../../..')) 12 | # from utils.o3d_tools import * 13 | from datasets.triplets.pointcloud_dataset import * 14 | 15 | class MulRanDataset(PointCloudDataset): 16 | r""" 17 | Generate single pointcloud frame from MulRan dataset. 18 | """ 19 | 20 | def __init__(self, 21 | phase, 22 | random_rotation=False, 23 | random_occlusion=False, 24 | random_scale=False, 25 | config=None): 26 | 27 | self.root = root = config["mulran_dir"] 28 | 29 | PointCloudDataset.__init__( 30 | self, phase, random_rotation, random_occlusion, random_scale, config) 31 | 32 | logging.info("Initializing MulRanDataset") 33 | logging.info(f"Loading the subset {phase} from {root}") 34 | 35 | # sequences = config.mulran_data_split[phase] 36 | # for drive_id in sequences: 37 | # inames = self.get_all_scan_ids(drive_id) 38 | # for query_id, start_time in enumerate(inames): 39 | # self.files.append((drive_id, query_id)) 40 | 41 | def get_all_scan_ids(self, drive_id): 42 | sequence_path = self.root + drive_id + '/Downsample/' 43 | fnames = sorted(glob.glob(os.path.join(sequence_path, '*.bin'))) 44 | assert len( 45 | fnames) > 0, f"Make sure that the path {self.root} has drive id: {drive_id}" 46 | inames = [int(os.path.split(fname)[-1][:-4]) for fname in fnames] 47 | return inames 48 | 49 | def get_velodyne_fn(self, drive_id, query_id): 50 | sequence_path = self.root + drive_id + '/Downsample/' 51 | fname = sorted(glob.glob(os.path.join( 52 | sequence_path, '*.bin')))[query_id] 53 | return fname 54 | 55 | def get_pointcloud_tensor(self, drive_id, pc_id): 56 | fname = self.get_velodyne_fn(drive_id, pc_id) 57 | pc = np.fromfile(fname, dtype=np.float64).reshape(-1, 3) 58 | 59 | if(pc.shape[0] != 4096) or (pc.shape[1] != 3): 60 | print("Error in pointcloud shape", pc.shape) 61 | return np.array([]) 62 | return pc 63 | 64 | def __getitem__(self, idx): 65 | drive_id = self.files[idx][0] 66 | t0 = self.files[idx][1] 67 | xyz0_th = self.get_pointcloud_tensor(drive_id, t0) 68 | meta_info = {'drive': drive_id, 't0': t0} 69 | 70 | return (xyz0_th, 71 | meta_info) 72 | 73 | class MulRanTupleDataset(MulRanDataset): 74 | r""" 75 | Generate tuples (anchor, positives, negatives) using distance 76 | Optional other_neg for quadruplet loss. 77 | """ 78 | 79 | def __init__(self, 80 | phase, 81 | random_rotation=False, 82 | random_occlusion=False, 83 | random_scale=False, 84 | config=None): 85 | self.root = root = config['mulran_dir'] 86 | self.positives_per_query = config['TRAIN_POSITIVES_PER_QUERY'] 87 | self.negatives_per_query = config['TRAIN_NEGATIVES_PER_QUERY'] 88 | self.quadruplet = False 89 | if config["LOSS_FUNCTION"] == 'quadruplet': 90 | self.quadruplet = True 91 | 92 | MulRanDataset.__init__( 93 | self, phase, random_rotation, random_occlusion, random_scale, config) 94 | 95 | logging.info("Initializing MulRanTupleDataset") 96 | logging.info(f"Loading the subset {phase} from {root}") 97 | 98 | sequences = ['DCC/DCC_01', 'DCC/DCC_02', 99 | 'Riverside/Riverside_01', 'Riverside/Riverside_03'] 100 | tuple_dir = os.path.join(os.path.dirname( 101 | __file__), '../config1/mulran_tuples/') 102 | self.dict_3m = json.load(open(tuple_dir + 'positive_sequence_D-3_T-0.json', "r")) 103 | self.dict_20m = json.load( 104 | open(tuple_dir + 'positive_sequence_D-20_T-0.json', "r")) 105 | self.mulran_seq_lens = {"DCC/DCC_01": 5542, "DCC/DCC_02": 7561, "DCC/DCC_03": 7479, 106 | "KAIST/KAIST_01": 8226, "KAIST/KAIST_02": 8941, "KAIST/KAIST_03": 8629, 107 | "Sejong/Sejong_01": 28779, "Sejong/Sejong_02": 27494, "Sejong/Sejong_03": 27215, 108 | "Riverside/Riverside_01": 5537, "Riverside/Riverside_02": 8157, "Riverside/Riverside_03": 10476} 109 | 110 | for drive_id in sequences: 111 | sequence_path = self.root + drive_id + '/Downsample/' 112 | fnames = sorted(glob.glob(os.path.join(sequence_path, '*.bin'))) 113 | assert len( 114 | fnames) > 0, f"Make sure that the path {root} has data {drive_id}" 115 | inames = sorted([int(os.path.split(fname)[-1][:-4]) 116 | for fname in fnames]) 117 | 118 | for query_id, start_time in enumerate(inames): 119 | positives = self.get_positives(drive_id, query_id) 120 | negatives = self.get_negatives(drive_id, query_id) 121 | self.files.append((drive_id, query_id, positives, negatives)) 122 | 123 | def get_positives(self, sq, index): 124 | assert sq in self.dict_3m.keys(), f"Error: Sequence {sq} not in json." 125 | sq_1 = self.dict_3m[sq] 126 | if str(int(index)) in sq_1: 127 | positives = sq_1[str(int(index))] 128 | else: 129 | positives = [] 130 | return positives 131 | 132 | def get_negatives(self, sq, index): 133 | assert sq in self.dict_20m.keys(), f"Error: Sequence {sq} not in json." 134 | sq_2 = self.dict_20m[sq] 135 | all_ids = set(np.arange(self.mulran_seq_lens[sq])) 136 | neg_set_inv = sq_2[str(int(index))] 137 | neg_set = all_ids.difference(neg_set_inv) 138 | negatives = list(neg_set) 139 | if index in negatives: 140 | negatives.remove(index) 141 | return negatives 142 | 143 | def get_other_negative(self, drive_id, query_id, sel_positive_ids, sel_negative_ids): 144 | # Dissimillar to all pointclouds in triplet tuple. 145 | all_ids = range(self.mulran_seq_lens[str(drive_id)]) 146 | neighbour_ids = sel_positive_ids 147 | for neg in sel_negative_ids: 148 | neg_postives_files = self.get_positives(drive_id, neg) 149 | for pos in neg_postives_files: 150 | neighbour_ids.append(pos) 151 | possible_negs = list(set(all_ids) - set(neighbour_ids)) 152 | if query_id in possible_negs: 153 | possible_negs.remove(query_id) 154 | assert len( 155 | possible_negs) > 0, f"No other negatives for drive {drive_id} id {query_id}" 156 | other_neg_id = random.sample(possible_negs, 1) 157 | return other_neg_id[0] 158 | 159 | def __getitem__(self, idx): 160 | drive_id, query_id = self.files[idx][0], self.files[idx][1] 161 | positive_ids, negative_ids = self.files[idx][2], self.files[idx][3] 162 | 163 | sel_positive_ids = random.sample( 164 | positive_ids, self.positives_per_query) 165 | sel_negative_ids = random.sample( 166 | negative_ids, self.negatives_per_query) 167 | positives, negatives, other_neg = [], [], None 168 | 169 | query_th = self.get_pointcloud_tensor(drive_id, query_id) 170 | for sp_id in sel_positive_ids: 171 | positives.append(self.get_pointcloud_tensor(drive_id, sp_id)) 172 | for sn_id in sel_negative_ids: 173 | negatives.append(self.get_pointcloud_tensor(drive_id, sn_id)) 174 | 175 | meta_info = {'drive': drive_id, 'query_id': query_id} 176 | 177 | if not self.quadruplet: 178 | return (query_th, 179 | positives, 180 | negatives, 181 | meta_info) 182 | else: # For Quadruplet Loss 183 | other_neg_id = self.get_other_negative( 184 | drive_id, query_id, sel_positive_ids, sel_negative_ids) 185 | other_neg_th = self.get_pointcloud_tensor(drive_id, other_neg_id) 186 | return (query_th, 187 | positives, 188 | negatives, 189 | other_neg_th, 190 | meta_info) 191 | 192 | ##################################################################################### 193 | # Load poses 194 | ##################################################################################### 195 | 196 | 197 | def load_poses_from_csv(file_name): 198 | with open(file_name, newline='') as f: 199 | reader = csv.reader(f) 200 | data_poses = list(reader) 201 | 202 | transforms = [] 203 | positions = [] 204 | for cnt, line in enumerate(data_poses): 205 | line_f = [float(i) for i in line] 206 | P = np.vstack((np.reshape(line_f[1:], (3, 4)), [0, 0, 0, 1])) 207 | transforms.append(P) 208 | positions.append([P[0, 3], P[1, 3], P[2, 3]]) 209 | return np.asarray(transforms), np.asarray(positions) 210 | 211 | 212 | ##################################################################################### 213 | # Load timestamps 214 | ##################################################################################### 215 | 216 | 217 | def load_timestamps_csv(file_name): 218 | with open(file_name, newline='') as f: 219 | reader = csv.reader(f) 220 | data_poses = list(reader) 221 | data_poses_ts = np.asarray( 222 | [float(t)/1e9 for t in np.asarray(data_poses)[:, 0]]) 223 | return data_poses_ts 224 | 225 | 226 | -------------------------------------------------------------------------------- /datasets/triplets/pointcloud_dataset.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | import torch 4 | import numpy as np 5 | 6 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 7 | # from utils.o3d_tools import * 8 | 9 | class PointCloudDataset(torch.utils.data.Dataset): 10 | def __init__(self, 11 | phase, 12 | random_rotation=False, 13 | random_occlusion=False, 14 | random_scale=False, 15 | config=None): 16 | self.phase = phase 17 | self.files = [] 18 | 19 | self.random_scale = random_scale 20 | self.min_scale = config["min_scale"] 21 | self.max_scale = config["max_scale"] 22 | self.random_occlusion = random_occlusion 23 | self.random_rotation = random_rotation 24 | self.rotation_range = 360 25 | if random_rotation: 26 | print('***********************Dataloader initialized with Random Rotation. ') 27 | if random_occlusion: 28 | print('***********************Dataloader initialized with Random Occlusion. ') 29 | if random_scale: 30 | print('***********************Dataloader initialized with Random Scale. ') 31 | 32 | def random_rotate(self, xyzr, r_angle=360, is_random=True, add_noise=True, rand_tr=False): 33 | # If is_random = True: Rotate about z-axis by random angle upto 'r_angle'. 34 | # Else: Rotate about z-axis by fixed angle 'r_angle'. 35 | r_angle = (np.pi/180) * r_angle 36 | if is_random: 37 | r_angle = r_angle*np.random.uniform() 38 | cos_angle = np.cos(r_angle) 39 | sin_angle = np.sin(r_angle) 40 | rot_matrix = np.array([[cos_angle, -sin_angle, 0], 41 | [sin_angle, cos_angle, 0], 42 | [0, 0, 1]]) 43 | scan = xyzr[:, :3] 44 | int = xyzr[:, 3].reshape((-1, 1)) 45 | augmented_scan = np.dot(scan, rot_matrix) 46 | 47 | if add_noise: 48 | n_sigma = 0.01 # Add gaussian noise 49 | noise = np.clip(n_sigma * np.random.randn(* 50 | augmented_scan.shape), -0.03, 0.03) 51 | augmented_scan = augmented_scan + noise 52 | 53 | if rand_tr: 54 | tr_xy_max, tr_z_max = 1.5, 0.25 55 | tr_xy = np.clip(np.random.randn(1, 2), -tr_xy_max, tr_xy_max) 56 | tr_z = np.clip(0.1*np.random.randn(1, 1), -tr_z_max, tr_z_max) 57 | tr = np.hstack((tr_xy, tr_z)) 58 | augmented_scan = augmented_scan + tr 59 | 60 | augmented_scan = np.hstack((augmented_scan, int)) 61 | return augmented_scan.astype(np.float32) 62 | 63 | def occlude_scan(self, scan, angle=30): 64 | # Remove points within a sector of fixed angle (degrees) and random heading direction. 65 | thetas = (180/np.pi) * np.arctan2(scan[:, 1], scan[:, 0]) 66 | heading = (180-angle/2)*np.random.uniform(-1, 1) 67 | occ_scan = np.vstack( 68 | (scan[thetas < (heading - angle/2)], scan[thetas > (heading + angle/2)])) 69 | return occ_scan.astype(np.float32) 70 | 71 | def pnv_preprocessing(self, xyzr, l=25): 72 | ind = np.argwhere(xyzr[:, 0] <= l).reshape(-1) 73 | xyzr = xyzr[ind] 74 | ind = np.argwhere(xyzr[:, 0] >= -l).reshape(-1) 75 | xyzr = xyzr[ind] 76 | ind = np.argwhere(xyzr[:, 1] <= l).reshape(-1) 77 | xyzr = xyzr[ind] 78 | ind = np.argwhere(xyzr[:, 1] >= -l).reshape(-1) 79 | xyzr = xyzr[ind] 80 | ind = np.argwhere(xyzr[:, 2] <= l).reshape(-1) 81 | xyzr = xyzr[ind] 82 | ind = np.argwhere(xyzr[:, 2] >= -l).reshape(-1) 83 | xyzr = xyzr[ind] 84 | 85 | vox_sz = 0.3 86 | while len(xyzr) > 4096: 87 | xyzr = downsample_point_cloud(xyzr, vox_sz) 88 | vox_sz += 0.01 89 | 90 | if xyzr.shape[0] >= 4096: 91 | ind = np.random.choice(xyzr.shape[0], 4096, replace=False) 92 | xyzr = xyzr[ind, :] 93 | else: 94 | ind = np.random.choice(xyzr.shape[0], 4096, replace=True) 95 | xyzr = xyzr[ind, :] 96 | mean = np.mean(xyzr, axis=0) 97 | pc = xyzr - mean 98 | scale = np.max(abs(pc)) 99 | pc = pc/scale 100 | return pc 101 | 102 | def __len__(self): 103 | return len(self.files) 104 | 105 | 106 | class CollationFunctionFactory: 107 | def __init__(self, collation_type='default', voxel_size=0.05, num_points=80000): 108 | self.voxel_size = voxel_size 109 | self.num_points = num_points 110 | if collation_type == 'default': 111 | self.collation_fn = self.collate_default 112 | elif collation_type == 'tuple': 113 | self.collation_fn = self.collate_tuple 114 | elif collation_type == 'sparse_tuple': 115 | self.collation_fn = self.collate_sparse_tuple 116 | elif collation_type == 'reg_sparse_tuple': 117 | self.collation_fn = self.collate_reg_sparse_tuple 118 | elif collation_type == 'sparcify_list': 119 | self.collation_fn = self.sparcify_and_collate_list 120 | else: 121 | raise ValueError(f'collation_type {collation_type} not found') 122 | 123 | def __call__(self, list_data): 124 | return self.collation_fn(list_data) 125 | 126 | def collate_default(self, list_data): 127 | if len(list_data) > 1: 128 | return self.collate_tuple(list_data) 129 | else: 130 | return list_data 131 | 132 | def collate_tuple(self, list_data): 133 | # print('1',len(list_data)) 134 | # print('2',len(list_data[0])) 135 | # for batch_data in list_data: 136 | # contrastive_tuple = [] 137 | # for tuple_data in batch_data: 138 | # print(len(tuple_data)) 139 | 140 | outputs = [] 141 | for batch_data in list_data: 142 | contrastive_tuple = [] 143 | for tuple_data in batch_data: 144 | if isinstance(tuple_data, np.ndarray): 145 | 146 | contrastive_tuple.append(tuple_data) 147 | elif isinstance(tuple_data, list): 148 | contrastive_tuple.extend(tuple_data) 149 | print(len(contrastive_tuple)) 150 | outputs.append(sparse_collate(contrastive_tuple)) 151 | if len(outputs) == 1: 152 | return outputs[0] 153 | else: 154 | return outputs 155 | 156 | def collate_sparse_tuple(self, list_data): 157 | outputs = [] 158 | for tuple_data in list_data: 159 | contrastive_tuple = [] 160 | for name in tuple_data.keys(): 161 | if isinstance(tuple_data[name], SparseTensor): 162 | contrastive_tuple.append(tuple_data[name]) 163 | elif isinstance(tuple_data[name], (list, np.ndarray)): 164 | contrastive_tuple.extend(tuple_data[name]) 165 | outputs.append(sparse_collate(contrastive_tuple)) 166 | if len(outputs) == 1: 167 | return outputs[0] 168 | else: 169 | return outputs 170 | 171 | def collate_reg_sparse_tuple(self, list_data): 172 | outputs = [] 173 | for tuple_data in list_data: 174 | contrastive_tuple = [] 175 | meta_info = None 176 | for name in tuple_data.keys(): 177 | if isinstance(tuple_data[name], SparseTensor): 178 | contrastive_tuple.append(tuple_data[name]) 179 | elif isinstance(tuple_data[name], (list, np.ndarray)): 180 | contrastive_tuple.extend(tuple_data[name]) 181 | elif isinstance(tuple_data[name], dict): 182 | meta_info = tuple_data[name] 183 | outputs.append([sparse_collate(contrastive_tuple), meta_info]) 184 | if len(outputs) == 1: 185 | return outputs[0] 186 | else: 187 | return outputs 188 | 189 | def sparcify_and_collate_list(self, list_data): 190 | outputs = [] 191 | if isinstance(list_data, SparseTensor): 192 | return list_data 193 | else: 194 | # return outputs 195 | for xyzr in list_data: 196 | xyzr = xyzr[0] 197 | if not len(xyzr) > 0: 198 | continue 199 | pc_ = np.round(xyzr[:, :3] / self.voxel_size).astype(np.int32) 200 | pc_ -= pc_.min(0, keepdims=1) 201 | feat_ = xyzr 202 | 203 | _, inds, inverse_map = sparse_quantize(pc_, 204 | return_index=True, 205 | return_inverse=True) 206 | if len(inds) > self.num_points: 207 | inds = np.random.choice( 208 | inds, self.num_points, replace=False) 209 | 210 | pc = pc_[inds] 211 | feat = feat_[inds] 212 | outputs.append(SparseTensor(feat, pc)) 213 | return sparse_collate(outputs) 214 | -------------------------------------------------------------------------------- /datasets/triplets/samplers.py: -------------------------------------------------------------------------------- 1 | # Warsaw University of Technology 2 | 3 | import random 4 | import copy 5 | from torch.utils.data import Sampler 6 | 7 | from datasets.triplets.base_dataset import TrainingDataset 8 | 9 | VERBOSE = False 10 | 11 | 12 | class ListDict(object): 13 | def __init__(self, items=None): 14 | if items is not None: 15 | self.items = copy.deepcopy(items) 16 | self.item_to_position = {item: ndx for ndx, item in enumerate(items)} 17 | else: 18 | self.items = [] 19 | self.item_to_position = {} 20 | 21 | def add(self, item): 22 | if item in self.item_to_position: 23 | return 24 | self.items.append(item) 25 | self.item_to_position[item] = len(self.items)-1 26 | 27 | def remove(self, item): 28 | position = self.item_to_position.pop(item) 29 | last_item = self.items.pop() 30 | if position != len(self.items): 31 | self.items[position] = last_item 32 | self.item_to_position[last_item] = position 33 | 34 | def choose_random(self): 35 | return random.choice(self.items) 36 | 37 | def __contains__(self, item): 38 | return item in self.item_to_position 39 | 40 | def __iter__(self): 41 | return iter(self.items) 42 | 43 | def __len__(self): 44 | return len(self.items) 45 | 46 | 47 | class BatchSampler(Sampler): 48 | # Sampler returning list of indices to form a mini-batch 49 | # Samples elements in groups consisting of k=2 similar elements (positives) 50 | # Batch has the following structure: item1_1, ..., item1_k, item2_1, ... item2_k, itemn_1, ..., itemn_k 51 | def __init__(self, dataset: TrainingDataset, batch_size: int, batch_size_limit: int = None, 52 | batch_expansion_rate: float = None, max_batches: int = None): 53 | if batch_expansion_rate is not None: 54 | assert batch_expansion_rate > 1., 'batch_expansion_rate must be greater than 1' 55 | assert batch_size <= batch_size_limit, 'batch_size_limit must be greater or equal to batch_size' 56 | 57 | self.batch_size = batch_size 58 | self.batch_size_limit = batch_size_limit 59 | self.batch_expansion_rate = batch_expansion_rate 60 | self.max_batches = max_batches 61 | self.dataset = dataset 62 | self.k = 2 # Number of positive examples per group must be 2 63 | if self.batch_size < 2 * self.k: 64 | self.batch_size = 2 * self.k 65 | print('WARNING: Batch too small. Batch size increased to {}.'.format(self.batch_size)) 66 | 67 | self.batch_idx = [] # Index of elements in each batch (re-generated every epoch) 68 | 69 | # self.elems_ndx = list(self.dataset.files) # List of point cloud indexes 70 | self.elems_ndx = list(range(0, len(self.dataset.files))) 71 | def __iter__(self): 72 | # Re-generate batches every epoch 73 | self.generate_batches() 74 | for batch in self.batch_idx: 75 | yield batch 76 | 77 | def __len(self): 78 | return len(self.batch_idx) 79 | 80 | def expand_batch(self): 81 | if self.batch_expansion_rate is None: 82 | print('WARNING: batch_expansion_rate is None') 83 | return 84 | 85 | if self.batch_size >= self.batch_size_limit: 86 | return 87 | 88 | old_batch_size = self.batch_size 89 | self.batch_size = int(self.batch_size * self.batch_expansion_rate) 90 | self.batch_size = min(self.batch_size, self.batch_size_limit) 91 | print('=> Batch size increased from: {} to {}'.format(old_batch_size, self.batch_size)) 92 | 93 | def generate_batches(self): 94 | # Generate training/evaluation batches. 95 | # batch_idx holds indexes of elements in each batch as a list of lists 96 | self.batch_idx = [] 97 | 98 | unused_elements_ndx = ListDict(self.elems_ndx) 99 | current_batch = [] 100 | 101 | assert self.k == 2, 'sampler can sample only k=2 elements from the same class' 102 | 103 | while True: 104 | if len(current_batch) >= self.batch_size or len(unused_elements_ndx) == 0: 105 | # Flush out batch, when it has a desired size, or a smaller batch, when there's no more 106 | # elements to process 107 | if len(current_batch) >= 2*self.k: 108 | # Ensure there're at least two groups of similar elements, otherwise, it would not be possible 109 | # to find negative examples in the batch 110 | assert len(current_batch) % self.k == 0, 'Incorrect bach size: {}'.format(len(current_batch)) 111 | self.batch_idx.append(current_batch) 112 | current_batch = [] 113 | if (self.max_batches is not None) and (len(self.batch_idx) >= self.max_batches): 114 | break 115 | if len(unused_elements_ndx) == 0: 116 | break 117 | 118 | # Add k=2 similar elements to the batch 119 | selected_element = unused_elements_ndx.choose_random() 120 | unused_elements_ndx.remove(selected_element) 121 | selected_element_drive_id, selected_element_query_id = self.dataset.files[selected_element][0], self.dataset.files[selected_element][1] 122 | positives = self.dataset.get_positives(selected_element_drive_id, selected_element_query_id) 123 | if len(positives) == 0: 124 | # Broken dataset element without any positives 125 | continue 126 | 127 | unused_positives = [e for e in positives if e in unused_elements_ndx] 128 | # If there're unused elements similar to selected_element, sample from them 129 | # otherwise sample from all similar elements 130 | if len(unused_positives) > 0: 131 | second_positive = random.choice(unused_positives) 132 | unused_elements_ndx.remove(second_positive) 133 | else: 134 | second_positive = random.choice(list(positives)) 135 | 136 | current_batch += [selected_element, second_positive] 137 | 138 | for batch in self.batch_idx: 139 | assert len(batch) % self.k == 0, 'Incorrect bach size: {}'.format(len(batch)) 140 | 141 | -------------------------------------------------------------------------------- /eval.sh: -------------------------------------------------------------------------------- 1 | 2 | # Re-localization on Oxford RobotCar/Inhouse(uni/res/bus) 3 | # python eval/relocalization.py --dataset_folder /home/xy/xy/code/Oxford/data/benchmark_datasets \ 4 | # --save_dir test/ \ 5 | # --backbone ppt_laws \ 6 | # --fc_output_dim 256 \ 7 | # --groups_num 8 \ 8 | # --batch_size 32 \ 9 | # --normalize_embeddings \ 10 | # --input_channel 1 \ 11 | # --eval_sequence oxf \ 12 | # --evaluate_model checkpoint_epoch_11.pth 13 | 14 | 15 | 16 | 17 | # Loop Closure on MulRan / KITTI 18 | 19 | # eval_dataset TestMulRanDataset --eval_dataset TestKittiDataset \ 20 | # --mulran_eval_seq 'Riverside/Riverside_02' \ --kitti_eval_seq 8 \ 21 | # --skip_time 90 30 22 | 23 | # Mulran sequences 24 | # KAIST/KAIST_01 25 | # DCC/DCC_03 26 | # Riverside/Riverside_02 27 | 28 | # KITTI sequences 29 | # 0 2 5 6 7 30 | 31 | python eval/loop_closure.py \ 32 | --eval_dataset TestKittiDataset \ 33 | --kitti_eval_seq 8 \ 34 | --backbone ppt_laws \ 35 | --eval_mode laws \ 36 | --groups_num 8 \ 37 | --input_channel 1 \ 38 | --pooling netvlad \ 39 | --normalize_embeddings \ 40 | --F1_thresh_id 526 \ 41 | --checkpoint_name 'logs/test/checkpoint_epoch_11.pth' \ 42 | --eval_batch_size 1 \ 43 | --skip_time 30 44 | 45 | 46 | 47 | # python eval/indoor_localization.py --checkpoint_name 'logs/test/checkpoint_epoch_11.pth' \ 48 | # --backbone ppt_laws \ 49 | # --eval_mode laws \ 50 | # --groups_num 8 \ 51 | # --input_channel 3 \ 52 | # --use_rgb \ 53 | # --fc_output_dim 256 \ 54 | # --batch_size 32 \ 55 | # --normalize_embeddings 56 | 57 | -------------------------------------------------------------------------------- /eval/__pycache__/eval_config.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/eval/__pycache__/eval_config.cpython-38.pyc -------------------------------------------------------------------------------- /eval/__pycache__/eval_sequence.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/eval/__pycache__/eval_sequence.cpython-38.pyc -------------------------------------------------------------------------------- /eval/__pycache__/plot_loop.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/eval/__pycache__/plot_loop.cpython-38.pyc -------------------------------------------------------------------------------- /eval/eval_config.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | arg_lists = [] 4 | parser = argparse.ArgumentParser() 5 | 6 | 7 | def add_argument_group(name): 8 | arg = parser.add_argument_group(name) 9 | arg_lists.append(arg) 10 | return arg 11 | 12 | 13 | def str2bool(v): 14 | return v.lower() in ('true', '1') 15 | 16 | 17 | # Evaluation 18 | eval_arg = add_argument_group('Eval') 19 | eval_arg.add_argument('--eval_pipeline', type=str, default='LOGG3D') 20 | eval_arg.add_argument('--kitti_eval_seq', type=int, default=8) 21 | eval_arg.add_argument('--mulran_eval_seq', type=str, 22 | default='Riverside/Riverside_02') 23 | eval_arg.add_argument('--wild_eval_seq', type=str, 24 | default='Karawatha/04') 25 | eval_arg.add_argument('--checkpoint_name', type=str, 26 | default='/kitti_10cm_loo/2021-09-14_20-28-22_3n24h_Kitti_v10_q29_10s8_263169.pth') 27 | eval_arg.add_argument('--eval_batch_size', type=int, default=1) 28 | eval_arg.add_argument('--test_num_workers', type=int, default=1) 29 | eval_arg.add_argument("--eval_random_rotation", type=str2bool, 30 | default=False, help="If random rotation. ") 31 | eval_arg.add_argument("--eval_random_occlusion", type=str2bool, 32 | default=False, help="If random occlusion. ") 33 | 34 | eval_arg.add_argument("--revisit_criteria", default=3, 35 | type=float, help="in meters") 36 | eval_arg.add_argument("--not_revisit_criteria", 37 | default=20, type=float, help="in meters") 38 | eval_arg.add_argument("--skip_time", default=30, type=float, help="in seconds") 39 | eval_arg.add_argument("--cd_thresh_min", default=0.001, 40 | type=float, help="Thresholds on cosine-distance to top-1.") 41 | eval_arg.add_argument("--cd_thresh_max", default=1.0, 42 | type=float, help="Thresholds on cosine-distance to top-1.") 43 | eval_arg.add_argument("--num_thresholds", default=1000, type=int, 44 | help="Number of thresholds. Number of points on PR curve.") 45 | eval_arg.add_argument("--radius", default=0.1, type=float, 46 | help="waiting to be completed") 47 | eval_arg.add_argument("--num_samples", default=64, type=int, 48 | help="waiting to be completed.") 49 | eval_arg.add_argument("--local_dim", default=256, type=int, 50 | help="waiting to be completed.") 51 | eval_arg.add_argument("--num_clusters", default=512, type=int, 52 | help="waiting to be completed.") 53 | 54 | # Minkloc parameters 55 | model_arg = add_argument_group('MinkModel') 56 | model_arg.add_argument("--planes", type=list, default=[64,128,64,32], help="_") 57 | model_arg.add_argument("--layers", type=list, default=[1,1,1,1], help="_") 58 | model_arg.add_argument("--num_top_down", type=int, default=2, help="_") 59 | model_arg.add_argument("--conv0_kernel_size", type=int, default=5, help="_") 60 | model_arg.add_argument("--feature_size", type=int, default=256, help="_") 61 | model_arg.add_argument("--output_dim", type=int, default=256, help="_") 62 | model_arg.add_argument("--block", type=str, default='ECABasicBlock', help="_") 63 | model_arg.add_argument("--pooling", type=str, default='netvlad', help="_") 64 | model_arg.add_argument("--coordinates", type=str, default='cartesian', help="_") 65 | model_arg.add_argument("--quantization_step", type=float, default=0.01, help="_") 66 | model_arg.add_argument("--normalize_embeddings", action='store_true') 67 | # Dataset specific configurations 68 | data_arg = add_argument_group('Data') 69 | # KittiDataset #MulRanDataset 70 | data_arg.add_argument('--eval_dataset', type=str, default='KittiDataset') 71 | data_arg.add_argument('--collation_type', type=str, 72 | default='default') # default#sparcify_list 73 | data_arg.add_argument("--eval_save_descriptors", type=str2bool, default=False) 74 | data_arg.add_argument("--eval_save_counts", type=str2bool, default=False) 75 | data_arg.add_argument("--gt_overlap", type=str2bool, default=False) 76 | data_arg.add_argument('--num_points', type=int, default=80000) 77 | data_arg.add_argument('--voxel_size', type=float, default=0.10) 78 | data_arg.add_argument("--gp_rem", type=str2bool, 79 | default=False, help="Remove ground plane.") 80 | data_arg.add_argument('--eval_feature_distance', type=str, 81 | default='cosine') # cosine#euclidean 82 | data_arg.add_argument("--pnv_preprocessing", type=str2bool, 83 | default=False, help="Preprocessing in dataloader for PNV.") 84 | 85 | data_arg.add_argument('--kitti_dir', type=str, default='/home/xy/xy/code/SemanticKitti/dataset/', 86 | help="Path to the KITTI odometry dataset") 87 | data_arg.add_argument('--kitti_data_split', type=dict, default={ 88 | 'train': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 89 | 'val': [], 90 | 'test': [0] 91 | }) 92 | 93 | data_arg.add_argument('--mulran_dir', type=str, 94 | default='/home/xy/xy/code/Mulran/', help="Path to the MulRan dataset") 95 | data_arg.add_argument("--mulran_normalize_intensity", type=str2bool, 96 | default=False, help="Normalize intensity return.") 97 | data_arg.add_argument('--mulran_data_split', type=dict, default={ 98 | 'train': ['DCC/DCC_01', 'DCC/DCC_02', 99 | 'Riverside/Riverside_01', 'Riverside/Riverside_03'], 100 | 'val': [], 101 | 'test': ['KAIST/KAIST_01'] 102 | }) 103 | data_arg.add_argument('--wild_dir', type=str, 104 | default='/home/xy/xy/Datasets/Wild-Places-master/datastore/kni101/wild-places/', help="Path to the Wild Place dataset") 105 | data_arg.add_argument('--wild_data_split', type=dict, default={ 106 | 'train': ['Karawatha/01', 'Karawatha/02', 107 | 'Venman/01', 'Venman/02'], 108 | 'val': [], 109 | 'test': ['Karawatha/03', 'Karawatha/04', 110 | 'Venman/03', 'Venman/04'] 111 | }) 112 | 113 | # data_arg.add_argument('--scannetpr_dir', type=str, 114 | # default='/root/siton-tmp/data/ScannetPR', help="Path to the MulRan dataset") 115 | data_arg.add_argument('--scannetpr_dir', type=str, 116 | default='/home/xy/xy/code/Data/ScannetPR', help="Path to the MulRan dataset") 117 | parser.add_argument("--use_rgb", action="store_true", 118 | help="use Automatic Mixed Precision") 119 | # Data loader configs 120 | data_arg.add_argument('--train_phase', type=str, default="train") 121 | data_arg.add_argument('--F1_thresh_id', type=int, default=100) 122 | data_arg.add_argument('--val_phase', type=str, default="val") 123 | data_arg.add_argument('--test_phase', type=str, default="test") 124 | data_arg.add_argument('--use_random_rotation', type=str2bool, default=False) 125 | data_arg.add_argument('--rotation_range', type=float, default=360) 126 | data_arg.add_argument('--use_random_occlusion', type=str2bool, default=False) 127 | data_arg.add_argument('--occlusion_angle', type=float, default=30) 128 | data_arg.add_argument('--use_random_scale', type=str2bool, default=False) 129 | data_arg.add_argument('--min_scale', type=float, default=0.8) 130 | data_arg.add_argument('--max_scale', type=float, default=1.2) 131 | data_arg.add_argument("--backbone", type=str, default="3d", 132 | choices=["ppt","ppt_laws","pnv","pnv_laws","mink","mink_laws"]) 133 | data_arg.add_argument("--batch_size", type=int, default=32) 134 | data_arg.add_argument("--groups_num", type=int, default=4) 135 | data_arg.add_argument("--fc_output_dim", type=int, default=256) 136 | data_arg.add_argument("--eval_mode", type=str, default='cp') 137 | data_arg.add_argument("--input_channel", type=int, default=3) 138 | data_arg.add_argument("--topN", type=int, default=5) 139 | # # PPT config 140 | data_arg.add_argument("--FEATURE_SIZE", type=list, default=[256,256,256,256], 141 | help="path of the folder with train/val/test sets") 142 | data_arg.add_argument("--MAX_SAMPLES", type=list, default=[64,256,1024,4096], 143 | help="__Undetermined") 144 | data_arg.add_argument("--CLUSTER_SIZE", type=list, default=[1,4,16,64], 145 | help="__Undetermined") 146 | data_arg.add_argument("--OUTPUT_DIM", type=list, default=[256,256,256,256], 147 | help="__Undetermined") 148 | data_arg.add_argument("--SAMPLING", type=list, default=[1024,256, 64,16], 149 | help="__Undetermined") 150 | data_arg.add_argument("--KNN", type=list, default=[20,20,20,20], 151 | help="__Undetermined") 152 | data_arg.add_argument("--GROUP", type=int, default=8, 153 | help="__Undetermined") 154 | data_arg.add_argument("--AGGREGATION", type=str, default='spvlad', 155 | help="__Undetermined") 156 | data_arg.add_argument("--GATING", type=bool, default= True, 157 | help="__Undetermined") 158 | 159 | 160 | def get_config_eval(): 161 | args = parser.parse_args() 162 | return args 163 | 164 | 165 | if __name__ == "__main__": 166 | cfg = get_config_eval() 167 | dconfig = vars(cfg) 168 | print(dconfig) 169 | -------------------------------------------------------------------------------- /eval/loop_closure.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import torch 4 | import logging 5 | sys.path.append(os.path.join(os.path.dirname(__file__), '../')) 6 | from eval.eval_sequence import * 7 | from models.model_factory import model_factory 8 | 9 | ch = logging.StreamHandler(sys.stdout) 10 | logging.getLogger().setLevel(logging.INFO) 11 | logging.basicConfig(format='%(asctime)s %(message)s', 12 | datefmt='%m/%d %H:%M:%S', 13 | handlers=[ch]) 14 | logging.basicConfig(level=logging.INFO, format="") 15 | 16 | 17 | def evaluate_checkpoint(model, save_path, cfg): 18 | # try: 19 | checkpoint = torch.load(save_path) # ,map_location='cuda:0') 20 | 21 | # try: 22 | # model.load_state_dict(checkpoint['model_state_dict']) 23 | # print('Use pretrain model') 24 | # except: 25 | # model.load_state_dict(checkpoint) 26 | # print('Retrieval model') 27 | 28 | model.load_state_dict(checkpoint['state_dict']) 29 | print('Use pretrain model') 30 | 31 | model = model.cuda() 32 | model.eval() 33 | return evaluate_sequence_reg(model, cfg) 34 | 35 | 36 | def evaluate_checkpoint_laws(model, save_path, cfg): 37 | # try: 38 | print('Use pretrain model') 39 | state = torch.load(save_path) # ,map_location='cuda:0') 40 | epoch = state['epoch'] + 1 41 | #### Updating 42 | for i in range(epoch): 43 | if i < cfg.groups_num: 44 | model.update_aggregators() 45 | print(model) 46 | #### Loading 47 | model.load_state_dict(state['model_state_dict']) 48 | model = model.cuda() 49 | model.eval() 50 | return evaluate_sequence_reg(model, cfg) 51 | 52 | 53 | if __name__ == "__main__": 54 | 55 | from eval.eval_config import get_config_eval 56 | cfg = get_config_eval() 57 | # Get model 58 | model = model_factory(cfg) 59 | model = model.to("cuda").eval() 60 | # model = DetectAndVLAD.get_model(args=cfg, radius = cfg.radius, num_samples =cfg.num_samples, feature_dim=cfg.local_dim, 61 | # num_clusters=cfg.num_clusters, batch_size = cfg.eval_batch_size) 62 | save_path = cfg.checkpoint_name 63 | print('Loading checkpoint from: ', save_path) 64 | logging.info('\n' + ' '.join([sys.executable] + sys.argv)) 65 | 66 | if cfg.eval_mode=='laws': 67 | eval_F1_max = evaluate_checkpoint_laws(model, save_path, cfg) 68 | elif cfg.eval_mode=='cp': 69 | eval_F1_max = evaluate_checkpoint(model, save_path, cfg) 70 | 71 | 72 | 73 | logging.info( 74 | '\n' + '******************* Evaluation Complete *******************') 75 | logging.info('Checkpoint Name: ' + str(cfg.checkpoint_name)) 76 | if 'Kitti' in cfg.eval_dataset: 77 | logging.info('Evaluated Sequence: ' + str(cfg.kitti_eval_seq)) 78 | elif 'MulRan' in cfg.eval_dataset: 79 | logging.info('Evaluated Sequence: ' + str(cfg.mulran_eval_seq)) 80 | logging.info('F1 Max: ' + str(eval_F1_max)) 81 | -------------------------------------------------------------------------------- /libs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/libs/__init__.py -------------------------------------------------------------------------------- /libs/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/libs/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /libs/pointops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/libs/pointops/__init__.py -------------------------------------------------------------------------------- /libs/pointops/functions/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/libs/pointops/functions/__init__.py -------------------------------------------------------------------------------- /libs/pointops/functions/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/libs/pointops/functions/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /libs/pointops/functions/__pycache__/pointops.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/libs/pointops/functions/__pycache__/pointops.cpython-36.pyc -------------------------------------------------------------------------------- /libs/pointops/setup.py: -------------------------------------------------------------------------------- 1 | #python3 setup.py install 2 | from setuptools import setup 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | 5 | setup( 6 | name='pointops', 7 | ext_modules=[ 8 | CUDAExtension('pointops_cuda', [ 9 | 'src/pointops_api.cpp', 10 | 'src/ballquery/ballquery_cuda.cpp', 11 | 'src/ballquery/ballquery_cuda_kernel.cu', 12 | 'src/knnquery/knnquery_cuda.cpp', 13 | 'src/knnquery/knnquery_cuda_kernel.cu', 14 | 'src/grouping/grouping_cuda.cpp', 15 | 'src/grouping/grouping_cuda_kernel.cu', 16 | 'src/grouping_int/grouping_int_cuda.cpp', 17 | 'src/grouping_int/grouping_int_cuda_kernel.cu', 18 | 'src/interpolation/interpolation_cuda.cpp', 19 | 'src/interpolation/interpolation_cuda_kernel.cu', 20 | 'src/sampling/sampling_cuda.cpp', 21 | 'src/sampling/sampling_cuda_kernel.cu', 22 | 23 | 'src/labelstat/labelstat_cuda.cpp', 24 | 'src/labelstat/labelstat_cuda_kernel.cu', 25 | 26 | 'src/featuredistribute/featuredistribute_cuda.cpp', 27 | 'src/featuredistribute/featuredistribute_cuda_kernel.cu' 28 | ], 29 | extra_compile_args={'cxx': ['-g'], 30 | 'nvcc': ['-O2']}) 31 | ], 32 | cmdclass={'build_ext': BuildExtension}) 33 | -------------------------------------------------------------------------------- /libs/pointops/src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/libs/pointops/src/__init__.py -------------------------------------------------------------------------------- /libs/pointops/src/ballquery/ballquery_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "ballquery_cuda_kernel.h" 7 | 8 | extern THCState *state; 9 | 10 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 11 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 12 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 13 | 14 | void ballquery_cuda(int b, int n, int m, float radius, int nsample, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) 15 | { 16 | const float *new_xyz = new_xyz_tensor.data(); 17 | const float *xyz = xyz_tensor.data(); 18 | int *idx = idx_tensor.data(); 19 | 20 | ballquery_cuda_launcher(b, n, m, radius, nsample, new_xyz, xyz, idx); 21 | } 22 | 23 | 24 | void ballquery_cuda_fast(int b, int n, int m, float radius, int nsample, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) 25 | { 26 | CHECK_INPUT(new_xyz_tensor); 27 | CHECK_INPUT(xyz_tensor); 28 | 29 | const float *new_xyz = new_xyz_tensor.data(); 30 | const float *xyz = xyz_tensor.data(); 31 | int *idx = idx_tensor.data(); 32 | 33 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 34 | 35 | ballquery_cuda_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx, stream); 36 | } 37 | -------------------------------------------------------------------------------- /libs/pointops/src/ballquery/ballquery_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "ballquery_cuda_kernel.h" 3 | 4 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 5 | // output: idx(b, m, nsample) 6 | __global__ void ballquery_cuda_kernel(int b, int n, int m, float radius, int nsample, const float *new_xyz, const float *xyz, int *idx) 7 | { 8 | int batch_index = blockIdx.x; 9 | xyz += batch_index * n * 3; 10 | new_xyz += batch_index * m * 3; 11 | idx += m * nsample * batch_index; 12 | int index = threadIdx.x; 13 | int stride = blockDim.x; 14 | 15 | float radius2 = radius * radius; 16 | for (int j = index; j < m; j += stride) 17 | { 18 | float new_x = new_xyz[j * 3 + 0]; 19 | float new_y = new_xyz[j * 3 + 1]; 20 | float new_z = new_xyz[j * 3 + 2]; 21 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) 22 | { 23 | float x = xyz[k * 3 + 0]; 24 | float y = xyz[k * 3 + 1]; 25 | float z = xyz[k * 3 + 2]; 26 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 27 | if (d2 < radius2) 28 | { 29 | if (cnt == 0) 30 | { 31 | for (int l = 0; l < nsample; ++l) 32 | idx[j * nsample + l] = k; 33 | } 34 | idx[j * nsample + cnt] = k; 35 | ++cnt; 36 | } 37 | } 38 | } 39 | } 40 | 41 | void ballquery_cuda_launcher(int b, int n, int m, float radius, int nsample, const float *new_xyz, const float *xyz, int *idx) 42 | { 43 | ballquery_cuda_kernel<<>>(b, n, m, radius, nsample, new_xyz, xyz, idx); 44 | } 45 | 46 | 47 | __global__ void ballquery_cuda_kernel_fast(int b, int n, int m, float radius, int nsample, const float *__restrict__ new_xyz, const float *__restrict__ xyz, int *__restrict__ idx) { 48 | int bs_idx = blockIdx.y; 49 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 50 | if (bs_idx >= b || pt_idx >= m) return; 51 | 52 | new_xyz += bs_idx * m * 3 + pt_idx * 3; 53 | xyz += bs_idx * n * 3; 54 | idx += bs_idx * m * nsample + pt_idx * nsample; 55 | 56 | float radius2 = radius * radius; 57 | float new_x = new_xyz[0]; 58 | float new_y = new_xyz[1]; 59 | float new_z = new_xyz[2]; 60 | 61 | int cnt = 0; 62 | for (int k = 0; k < n; ++k) { 63 | float x = xyz[k * 3 + 0]; 64 | float y = xyz[k * 3 + 1]; 65 | float z = xyz[k * 3 + 2]; 66 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 67 | if (d2 < radius2){ 68 | if (cnt == 0){ 69 | for (int l = 0; l < nsample; ++l) { 70 | idx[l] = k; 71 | } 72 | } 73 | idx[cnt] = k; 74 | ++cnt; 75 | if (cnt >= nsample){ 76 | break; 77 | } 78 | } 79 | } 80 | } 81 | 82 | 83 | void ballquery_cuda_launcher_fast(int b, int n, int m, float radius, int nsample, const float *new_xyz, const float *xyz, int *idx, cudaStream_t stream) { 84 | // param new_xyz: (B, m, 3) 85 | // param xyz: (B, n, 3) 86 | // param idx: (B, m, nsample) 87 | 88 | cudaError_t err; 89 | 90 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 91 | dim3 threads(THREADS_PER_BLOCK); 92 | 93 | ballquery_cuda_kernel_fast<<>>(b, n, m, radius, nsample, new_xyz, xyz, idx); 94 | // cudaDeviceSynchronize(); // for using printf in kernel function 95 | 96 | err = cudaGetLastError(); 97 | if (cudaSuccess != err) { 98 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 99 | exit(-1); 100 | } 101 | } 102 | -------------------------------------------------------------------------------- /libs/pointops/src/ballquery/ballquery_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _BALLQUERY_CUDA_KERNEL 2 | #define _BALLQUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void ballquery_cuda(int b, int n, int m, float radius, int nsample, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor); 8 | 9 | void ballquery_cuda_fast(int b, int n, int m, float radius, int nsample, at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor); 10 | 11 | #ifdef __cplusplus 12 | extern "C" { 13 | #endif 14 | 15 | void ballquery_cuda_launcher(int b, int n, int m, float radius, int nsample, const float *xyz, const float *new_xyz, int *idx); 16 | 17 | void ballquery_cuda_launcher_fast(int b, int n, int m, float radius, int nsample, const float *new_xyz, const float *xyz, int *idx, cudaStream_t stream); 18 | 19 | #ifdef __cplusplus 20 | } 21 | #endif 22 | 23 | #endif 24 | -------------------------------------------------------------------------------- /libs/pointops/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | 6 | #define TOTAL_THREADS 1024 7 | 8 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 9 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 10 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 11 | 12 | #define THREADS_PER_BLOCK 256 13 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) 14 | 15 | inline int opt_n_threads(int work_size) { 16 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 17 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 18 | } 19 | 20 | inline dim3 opt_block_config(int x, int y) { 21 | const int x_threads = opt_n_threads(x); 22 | const int y_threads = max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 23 | dim3 block_config(x_threads, y_threads, 1); 24 | return block_config; 25 | } 26 | 27 | #endif -------------------------------------------------------------------------------- /libs/pointops/src/featuredistribute/featuredistribute_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "featuredistribute_cuda_kernel.h" 7 | 8 | extern THCState *state; 9 | 10 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 11 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 12 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 13 | 14 | 15 | void featuredistribute_cuda(int b, int n, int m, at::Tensor max_xyz_tensor, at::Tensor xyz_tensor, at::Tensor distribute_idx_tensor) 16 | { 17 | CHECK_INPUT(max_xyz_tensor); 18 | CHECK_INPUT(xyz_tensor); 19 | 20 | const float *max_xyz = max_xyz_tensor.data(); 21 | const float *xyz = xyz_tensor.data(); 22 | int *distribute_idx = distribute_idx_tensor.data(); 23 | 24 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 25 | 26 | featuredistribute_cuda_launcher(b, n, m, max_xyz, xyz, distribute_idx, stream); 27 | } 28 | 29 | 30 | void featuregather_forward_cuda(int b, int n, int m, int c, at::Tensor max_feature_tensor, at::Tensor distribute_idx_tensor, at::Tensor distribute_feature_tensor) 31 | { 32 | CHECK_INPUT(max_feature_tensor); 33 | CHECK_INPUT(distribute_idx_tensor); 34 | 35 | const float *max_feature = max_feature_tensor.data(); 36 | const int *distribute_idx = distribute_idx_tensor.data(); 37 | float *distribute_feature = distribute_feature_tensor.data(); 38 | 39 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 40 | 41 | featuregather_forward_cuda_launcher(b, n, m, c, max_feature, distribute_idx, distribute_feature, stream); 42 | } 43 | 44 | 45 | void featuregather_backward_cuda(int b, int n, int m, int c, at::Tensor grad_distribute_feature_tensor, at::Tensor distribute_idx_tensor, at::Tensor grad_max_feature_tensor) 46 | { 47 | CHECK_INPUT(grad_distribute_feature_tensor); 48 | CHECK_INPUT(distribute_idx_tensor); 49 | 50 | const float *grad_distribute_feature = grad_distribute_feature_tensor.data(); 51 | const int *distribute_idx = distribute_idx_tensor.data(); 52 | float *grad_max_feature = grad_max_feature_tensor.data(); 53 | 54 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 55 | 56 | featuregather_backward_cuda_launcher(b, n, m, c, grad_distribute_feature, distribute_idx, grad_max_feature, stream); 57 | } -------------------------------------------------------------------------------- /libs/pointops/src/featuredistribute/featuredistribute_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "featuredistribute_cuda_kernel.h" 3 | 4 | __global__ void featuredistribute_cuda_kernel(int b, int n, int m, const float *max_xyz, const float *xyz, int *distribute_idx) { 5 | int bs_idx = blockIdx.y; 6 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 7 | if (bs_idx >= b || pt_idx >= m) return; 8 | 9 | max_xyz += bs_idx * n * 3; 10 | xyz += bs_idx * m * 3 + pt_idx * 3; 11 | distribute_idx += bs_idx * m + pt_idx; 12 | 13 | float x = xyz[0]; 14 | float y = xyz[1]; 15 | float z = xyz[2]; 16 | 17 | float min_dist2 = 100000; 18 | int min_dist_idx = -1; 19 | for (int k = 0; k < n; ++k) { 20 | float max_x = max_xyz[k * 3 + 0]; 21 | float max_y = max_xyz[k * 3 + 1]; 22 | float max_z = max_xyz[k * 3 + 2]; 23 | float d2 = (max_x - x) * (max_x - x) + (max_y - y) * (max_y - y) + (max_z - z) * (max_z - z); 24 | if (d2 < min_dist2){ 25 | min_dist_idx = k; 26 | min_dist2 = d2; 27 | } 28 | } 29 | distribute_idx[0] = min_dist_idx; 30 | } 31 | 32 | 33 | void featuredistribute_cuda_launcher(int b, int n, int m, const float *max_xyz, const float *xyz, int *distribute_idx, cudaStream_t stream) { 34 | // param max_xyz: (b, n, 3) 35 | // param xyz: (b, m, 3) 36 | // return distribute_idx: (b, m) 37 | 38 | cudaError_t err; 39 | 40 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 41 | dim3 threads(THREADS_PER_BLOCK); 42 | 43 | featuredistribute_cuda_kernel<<>>(b, n, m, max_xyz, xyz, distribute_idx); 44 | // cudaDeviceSynchronize(); // for using printf in kernel function 45 | 46 | err = cudaGetLastError(); 47 | if (cudaSuccess != err) { 48 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 49 | exit(-1); 50 | } 51 | } 52 | 53 | __global__ void featuregather_forward_cuda_kernel(int b, int n, int m, int c, const float *max_feature, const int *distribute_idx, float *distribute_feature) { 54 | int bs_idx = blockIdx.z; 55 | int c_idx = blockIdx.y; 56 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 57 | if (bs_idx >= b || c_idx >= c || pt_idx >= m) return; 58 | 59 | max_feature += bs_idx * c * n + c_idx * n; 60 | distribute_idx += bs_idx * m + pt_idx; 61 | distribute_feature += bs_idx * c * m + c_idx * m + pt_idx; 62 | 63 | int idx = distribute_idx[0]; 64 | distribute_feature[0] = max_feature[idx]; 65 | } 66 | 67 | 68 | void featuregather_forward_cuda_launcher(int b, int n, int m, int c, const float *max_feature, const int *distribute_idx, float *distribute_feature, cudaStream_t stream){ 69 | // param max_feature: (b, c, n) 70 | // param distribute_idx: (b, m) 71 | // return distribute_feature: (b, c, m) 72 | 73 | cudaError_t err; 74 | 75 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) 76 | dim3 threads(THREADS_PER_BLOCK); 77 | 78 | featuregather_forward_cuda_kernel<<>>(b, n, m, c, max_feature, distribute_idx, distribute_feature); 79 | // cudaDeviceSynchronize(); // for using printf in kernel function 80 | 81 | err = cudaGetLastError(); 82 | if (cudaSuccess != err) { 83 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 84 | exit(-1); 85 | } 86 | } 87 | 88 | 89 | __global__ void featuregather_backward_cuda_kernel(int b, int n, int m, int c, const float *grad_distribute_feature, const int *distribute_idx, float *grad_max_feature){ 90 | int bs_idx = blockIdx.z; 91 | int c_idx = blockIdx.y; 92 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 93 | if(bs_idx >= b || c_idx >= c || pt_idx >= m) return; 94 | 95 | grad_distribute_feature += bs_idx * c * m + c_idx * m + pt_idx; 96 | distribute_idx += bs_idx * m + pt_idx; 97 | grad_max_feature += bs_idx * c * n + c_idx * n; 98 | 99 | int idx = distribute_idx[0]; 100 | atomicAdd(grad_max_feature + idx, grad_distribute_feature[0]); 101 | } 102 | 103 | 104 | void featuregather_backward_cuda_launcher(int b, int n, int m, int c, const float *grad_distribute_feature, const int *distribute_idx, float *grad_max_feature, cudaStream_t stream){ 105 | // param grad_distribute_feature: (b, c, m) 106 | // param distribute_idx: (b, m) 107 | // return grad_max_feature: (b, c, n) 108 | 109 | cudaError_t err; 110 | 111 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) 112 | dim3 threads(THREADS_PER_BLOCK); 113 | 114 | featuregather_backward_cuda_kernel<<>>(b, n, m, c, grad_distribute_feature, distribute_idx, grad_max_feature); 115 | // cudaDeviceSynchronize(); // for using printf in kernel function 116 | 117 | err = cudaGetLastError(); 118 | if (cudaSuccess != err) { 119 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 120 | exit(-1); 121 | } 122 | } -------------------------------------------------------------------------------- /libs/pointops/src/featuredistribute/featuredistribute_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _FEATUREDISTRIBUTE_CUDA_KERNEL 2 | #define _FEATUREDISTRIBUTE_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void featuredistribute_cuda(int b, int n, int m, at::Tensor max_xyz_tensor, at::Tensor xyz_tensor, at::Tensor distribute_idx_tensor); 8 | void featuregather_forward_cuda(int b, int n, int m, int c, at::Tensor max_feature_tensor, at::Tensor distribute_idx_tensor, at::Tensor distribute_feature_tensor); 9 | void featuregather_backward_cuda(int b, int n, int m, int c, at::Tensor grad_distribute_feature_tensor, at::Tensor distribute_idx_tensor, at::Tensor grad_max_feature_tensor); 10 | 11 | #ifdef __cplusplus 12 | extern "C" { 13 | #endif 14 | 15 | void featuredistribute_cuda_launcher(int b, int n, int m, const float *max_xyz, const float *xyz, int *distribute_idx, cudaStream_t stream); 16 | void featuregather_forward_cuda_launcher(int b, int n, int m, int c, const float *max_feature, const int *distribute_idx, float *distribute_feature, cudaStream_t stream); 17 | void featuregather_backward_cuda_launcher(int b, int n, int m, int c, const float *grad_distribute_feature, const int *distribute_idx, float *grad_max_feature, cudaStream_t stream); 18 | 19 | #ifdef __cplusplus 20 | } 21 | #endif 22 | 23 | #endif 24 | -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "grouping_cuda_kernel.h" 7 | 8 | extern THCState *state; 9 | 10 | void grouping_forward_cuda(int b, int c, int n, int m, int nsample, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) 11 | { 12 | const float *points = points_tensor.data(); 13 | const int *idx = idx_tensor.data(); 14 | float *out = out_tensor.data(); 15 | grouping_forward_cuda_launcher(b, c, n, m, nsample, points, idx, out); 16 | } 17 | 18 | void grouping_backward_cuda(int b, int c, int n, int m, int nsample, at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) 19 | { 20 | float *grad_points = grad_points_tensor.data(); 21 | const int *idx = idx_tensor.data(); 22 | const float *grad_out = grad_out_tensor.data(); 23 | grouping_backward_cuda_launcher(b, c, n, m, nsample, grad_out, idx, grad_points); 24 | } 25 | 26 | void grouping_forward_cuda_fast(int b, int c, int n, int npoints, int nsample, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) { 27 | 28 | const float *points = points_tensor.data(); 29 | const int *idx = idx_tensor.data(); 30 | float *out = out_tensor.data(); 31 | grouping_forward_cuda_launcher_fast(b, c, n, npoints, nsample, points, idx, out); 32 | } -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "grouping_cuda_kernel.h" 3 | 4 | // input: points(b, c, n) idx(b, m, nsample) 5 | // output: out(b, c, m, nsample) 6 | __global__ void grouping_forward_cuda_kernel(int b, int c, int n, int m, int nsample, const float *points, const int *idx, float *out) 7 | { 8 | int batch_index = blockIdx.x; 9 | points += batch_index * n * c; 10 | idx += batch_index * m * nsample; 11 | out += batch_index * m * nsample * c; 12 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 13 | const int stride = blockDim.y * blockDim.x; 14 | for (int i = index; i < c * m; i += stride) 15 | { 16 | const int l = i / m; 17 | const int j = i % m; 18 | for (int k = 0; k < nsample; ++k) 19 | { 20 | int ii = idx[j * nsample + k]; 21 | out[(l * m + j) * nsample + k] = points[l * n + ii]; 22 | } 23 | } 24 | } 25 | 26 | // input: grad_out(b, c, m, nsample), idx(b, m, nsample) 27 | // output: grad_points(b, c, n) 28 | __global__ void grouping_backward_cuda_kernel(int b, int c, int n, int m, int nsample, const float *grad_out, const int *idx, float *grad_points) 29 | { 30 | int batch_index = blockIdx.x; 31 | grad_out += batch_index * m * nsample * c; 32 | idx += batch_index * m * nsample; 33 | grad_points += batch_index * n * c; 34 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 35 | const int stride = blockDim.y * blockDim.x; 36 | for (int i = index; i < c * m; i += stride) 37 | { 38 | const int l = i / m; 39 | const int j = i % m; 40 | for (int k = 0; k < nsample; ++k) 41 | { 42 | int ii = idx[j * nsample + k]; 43 | atomicAdd(grad_points + l * n + ii, grad_out[(l * m + j) * nsample + k]); 44 | } 45 | } 46 | } 47 | 48 | void grouping_forward_cuda_launcher(int b, int c, int n, int m, int nsample, const float *points, const int *idx, float *out) 49 | { 50 | grouping_forward_cuda_kernel<<>>(b, c, n, m, nsample, points, idx, out); 51 | } 52 | 53 | void grouping_backward_cuda_launcher(int b, int c, int n, int m, int nsample, const float *grad_out, const int *idx, float *grad_points) 54 | { 55 | grouping_backward_cuda_kernel<<>>(b, c, n, m, nsample, grad_out, idx, grad_points); 56 | } 57 | 58 | // input: points(b, c, n) idx(b, npoints, nsample) 59 | // output: out(b, c, npoints, nsample) 60 | __global__ void grouping_forward_cuda_kernel_fast(int b, int c, int n, int npoints, int nsample, const float *__restrict__ points, const int *__restrict__ idx, float *__restrict__ out) { 61 | int bs_idx = blockIdx.z; 62 | int c_idx = blockIdx.y; 63 | int index = blockIdx.x * blockDim.x + threadIdx.x; 64 | int pt_idx = index / nsample; 65 | if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return; 66 | 67 | int sample_idx = index % nsample; 68 | 69 | idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx; 70 | int in_idx = bs_idx * c * n + c_idx * n + idx[0]; 71 | int out_idx = bs_idx * c * npoints * nsample + c_idx * npoints * nsample + pt_idx * nsample + sample_idx; 72 | 73 | out[out_idx] = points[in_idx]; 74 | } 75 | 76 | // input: points(b, c, n) idx(b, npoints, nsample) 77 | // output: out(b, c, npoints, nsample) 78 | void grouping_forward_cuda_launcher_fast(int b, int c, int n, int npoints, int nsample, const float *points, const int *idx, float *out) { 79 | 80 | cudaError_t err; 81 | 82 | dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) 83 | dim3 threads(THREADS_PER_BLOCK); 84 | 85 | grouping_forward_cuda_kernel_fast<<>>(b, c, n, npoints, nsample, points, idx, out); 86 | // cudaDeviceSynchronize(); // for using printf in kernel function 87 | err = cudaGetLastError(); 88 | if (cudaSuccess != err) { 89 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 90 | exit(-1); 91 | } 92 | } 93 | 94 | 95 | -------------------------------------------------------------------------------- /libs/pointops/src/grouping/grouping_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _GROUPING_CUDA_KERNEL 2 | #define _GROUPING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void grouping_forward_cuda(int b, int c, int n, int m, int nsample, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out); 8 | void grouping_backward_cuda(int b, int c, int n, int m, int nsample, at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor); 9 | 10 | void grouping_forward_cuda_fast(int b, int c, int n, int npoints, int nsample, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor); 11 | 12 | #ifdef __cplusplus 13 | extern "C" { 14 | #endif 15 | 16 | void grouping_forward_cuda_launcher(int b, int c, int n, int m, int nsample, const float *points, const int *idx, float *out); 17 | void grouping_backward_cuda_launcher(int b, int c, int n, int m, int nsample, const float *grad_out, const int *idx, float *grad_points); 18 | 19 | void grouping_forward_cuda_launcher_fast(int b, int c, int n, int npoints, int nsample, const float *points, const int *idx, float *out); 20 | 21 | #ifdef __cplusplus 22 | } 23 | #endif 24 | #endif 25 | -------------------------------------------------------------------------------- /libs/pointops/src/grouping_int/grouping_int_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "grouping_int_cuda_kernel.h" 7 | 8 | extern THCState *state; 9 | 10 | void grouping_int_forward_cuda(int b, int c, int n, int m, int nsample, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) 11 | { 12 | const long int *points = points_tensor.data(); 13 | const int *idx = idx_tensor.data(); 14 | long int *out = out_tensor.data(); 15 | grouping_int_forward_cuda_launcher(b, c, n, m, nsample, points, idx, out); 16 | } 17 | 18 | void grouping_int_forward_cuda_fast(int b, int c, int n, int m, int nsample, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) 19 | { 20 | const long int *points = points_tensor.data(); 21 | const int *idx = idx_tensor.data(); 22 | long int *out = out_tensor.data(); 23 | grouping_int_forward_cuda_launcher_fast(b, c, n, m, nsample, points, idx, out); 24 | } -------------------------------------------------------------------------------- /libs/pointops/src/grouping_int/grouping_int_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "grouping_int_cuda_kernel.h" 3 | 4 | // input: points(b, c, n) idx(b, m, nsample) 5 | // output: out(b, c, m, nsample) 6 | __global__ void grouping_int_forward_cuda_kernel(int b, int c, int n, int m, int nsample, const long int *points, const int *idx, long int *out) 7 | { 8 | int batch_index = blockIdx.x; 9 | points += batch_index * n * c; 10 | idx += batch_index * m * nsample; 11 | out += batch_index * m * nsample * c; 12 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 13 | const int stride = blockDim.y * blockDim.x; 14 | for (int i = index; i < c * m; i += stride) 15 | { 16 | const int l = i / m; 17 | const int j = i % m; 18 | for (int k = 0; k < nsample; ++k) 19 | { 20 | int ii = idx[j * nsample + k]; 21 | out[(l * m + j) * nsample + k] = points[l * n + ii]; 22 | } 23 | } 24 | } 25 | 26 | 27 | void grouping_int_forward_cuda_launcher(int b, int c, int n, int m, int nsample, const long int *points, const int *idx, long int *out) 28 | { 29 | grouping_int_forward_cuda_kernel<<>>(b, c, n, m, nsample, points, idx, out); 30 | } 31 | 32 | 33 | __global__ void grouping_int_forward_cuda_kernel_fast(int b, int c, int n, int npoints, int nsample, const long int *__restrict__ points, const int *__restrict__ idx, long int *__restrict__ out) 34 | { 35 | int bs_idx = blockIdx.z; 36 | int c_idx = blockIdx.y; 37 | int index = blockIdx.x * blockDim.x + threadIdx.x; 38 | int pt_idx = index / nsample; 39 | if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return; 40 | 41 | int sample_idx = index % nsample; 42 | 43 | idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx; 44 | int in_idx = bs_idx * c * n + c_idx * n + idx[0]; 45 | int out_idx = bs_idx * c * npoints * nsample + c_idx * npoints * nsample + pt_idx * nsample + sample_idx; 46 | 47 | out[out_idx] = points[in_idx]; 48 | } 49 | 50 | 51 | void grouping_int_forward_cuda_launcher_fast(int b, int c, int n, int npoints, int nsample, const long int *points, const int *idx, long int *out) 52 | { 53 | cudaError_t err; 54 | 55 | dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) 56 | dim3 threads(THREADS_PER_BLOCK); 57 | 58 | grouping_int_forward_cuda_kernel_fast<<>>(b, c, n, npoints, nsample, points, idx, out); 59 | // cudaDeviceSynchronize(); // for using printf in kernel function 60 | err = cudaGetLastError(); 61 | if (cudaSuccess != err) { 62 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 63 | exit(-1); 64 | } 65 | } -------------------------------------------------------------------------------- /libs/pointops/src/grouping_int/grouping_int_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _GROUPING_INT_CUDA_KERNEL 2 | #define _GROUPING_INT_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void grouping_int_forward_cuda(int b, int c, int n, int m, int nsample, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out); 8 | 9 | void grouping_int_forward_cuda_fast(int b, int c, int n, int m, int nsample, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor); 10 | 11 | #ifdef __cplusplus 12 | extern "C" { 13 | #endif 14 | 15 | void grouping_int_forward_cuda_launcher(int b, int c, int n, int m, int nsample, const long int *points, const int *idx, long int *out); 16 | 17 | void grouping_int_forward_cuda_launcher_fast(int b, int c, int n, int npoints, int nsample, const long int *points, const int *idx, long int *out); 18 | 19 | #ifdef __cplusplus 20 | } 21 | #endif 22 | #endif 23 | -------------------------------------------------------------------------------- /libs/pointops/src/interpolation/interpolation_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "interpolation_cuda_kernel.h" 6 | 7 | extern THCState *state; 8 | 9 | void nearestneighbor_cuda(int b, int n, int m, at::Tensor unknown_tensor, at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor) 10 | { 11 | const float *unknown = unknown_tensor.data(); 12 | const float *known = known_tensor.data(); 13 | float *dist2 = dist2_tensor.data(); 14 | int *idx = idx_tensor.data(); 15 | nearestneighbor_cuda_launcher(b, n, m, unknown, known, dist2, idx); 16 | } 17 | 18 | void interpolation_forward_cuda(int b, int c, int m, int n, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor) 19 | { 20 | const float *points = points_tensor.data(); 21 | const float *weight = weight_tensor.data(); 22 | float *out = out_tensor.data(); 23 | const int *idx = idx_tensor.data(); 24 | interpolation_forward_cuda_launcher(b, c, m, n, points, idx, weight, out); 25 | } 26 | 27 | void interpolation_backward_cuda(int b, int c, int n, int m, at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_points_tensor) 28 | { 29 | const float *grad_out = grad_out_tensor.data(); 30 | const float *weight = weight_tensor.data(); 31 | float *grad_points = grad_points_tensor.data(); 32 | const int *idx = idx_tensor.data(); 33 | interpolation_backward_cuda_launcher(b, c, n, m, grad_out, idx, weight, grad_points); 34 | } 35 | 36 | void nearestneighbor_cuda_fast(int b, int n, int m, at::Tensor unknown_tensor, at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor) { 37 | const float *unknown = unknown_tensor.data(); 38 | const float *known = known_tensor.data(); 39 | float *dist2 = dist2_tensor.data(); 40 | int *idx = idx_tensor.data(); 41 | nearestneighbor_cuda_launcher_fast(b, n, m, unknown, known, dist2, idx); 42 | } 43 | 44 | void interpolation_forward_cuda_fast(int b, int c, int m, int n, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor) { 45 | 46 | const float *points = points_tensor.data(); 47 | const float *weight = weight_tensor.data(); 48 | float *out = out_tensor.data(); 49 | const int *idx = idx_tensor.data(); 50 | interpolation_forward_cuda_launcher_fast(b, c, m, n, points, idx, weight, out); 51 | } -------------------------------------------------------------------------------- /libs/pointops/src/interpolation/interpolation_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "interpolation_cuda_kernel.h" 3 | 4 | // input: unknown(b, n, 3) known(b, m, 3) 5 | // output: dist2(b, n, 3), idx(b, n, 3) 6 | __global__ void nearestneighbor_cuda_kernel(int b, int n, int m, const float *unknown, const float *known, float *dist2, int *idx) 7 | { 8 | int batch_index = blockIdx.x; 9 | unknown += batch_index * n * 3; 10 | known += batch_index * m * 3; 11 | dist2 += batch_index * n * 3; 12 | idx += batch_index * n * 3; 13 | 14 | int index = threadIdx.x; 15 | int stride = blockDim.x; 16 | for (int j = index; j < n; j += stride) 17 | { 18 | float ux = unknown[j * 3 + 0]; 19 | float uy = unknown[j * 3 + 1]; 20 | float uz = unknown[j * 3 + 2]; 21 | 22 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 23 | int besti1 = 0, besti2 = 0, besti3 = 0; 24 | for (int k = 0; k < m; ++k) 25 | { 26 | float x = known[k * 3 + 0]; 27 | float y = known[k * 3 + 1]; 28 | float z = known[k * 3 + 2]; 29 | float d = 30 | (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 31 | if (d < best1) 32 | { 33 | best3 = best2; 34 | besti3 = besti2; 35 | best2 = best1; 36 | besti2 = besti1; 37 | best1 = d; 38 | besti1 = k; 39 | } 40 | else if (d < best2) 41 | { 42 | best3 = best2; 43 | besti3 = besti2; 44 | best2 = d; 45 | besti2 = k; 46 | } 47 | else if (d < best3) 48 | { 49 | best3 = d; 50 | besti3 = k; 51 | } 52 | } 53 | dist2[j * 3 + 0] = best1; 54 | dist2[j * 3 + 1] = best2; 55 | dist2[j * 3 + 2] = best3; 56 | idx[j * 3 + 0] = besti1; 57 | idx[j * 3 + 1] = besti2; 58 | idx[j * 3 + 2] = besti3; 59 | } 60 | } 61 | 62 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 63 | // output: out(b, c, n) 64 | __global__ void interpolation_forward_cuda_kernel(int b, int c, int m, int n, const float *points, const int *idx, const float *weight, float *out) 65 | { 66 | int batch_index = blockIdx.x; 67 | points += batch_index * m * c; 68 | idx += batch_index * n * 3; 69 | weight += batch_index * n * 3; 70 | out += batch_index * n * c; 71 | 72 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 73 | const int stride = blockDim.y * blockDim.x; 74 | for (int i = index; i < c * n; i += stride) 75 | { 76 | const int l = i / n; 77 | const int j = i % n; 78 | float w1 = weight[j * 3 + 0]; 79 | float w2 = weight[j * 3 + 1]; 80 | float w3 = weight[j * 3 + 2]; 81 | int i1 = idx[j * 3 + 0]; 82 | int i2 = idx[j * 3 + 1]; 83 | int i3 = idx[j * 3 + 2]; 84 | out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + points[l * m + i3] * w3; 85 | } 86 | } 87 | 88 | // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) 89 | // output: grad_points(b, c, m) 90 | __global__ void interpolation_backward_cuda_kernel( int b, int c, int n, int m, const float *grad_out, const int *idx, const float *weight, float *grad_points) 91 | { 92 | int batch_index = blockIdx.x; 93 | grad_out += batch_index * n * c; 94 | idx += batch_index * n * 3; 95 | weight += batch_index * n * 3; 96 | grad_points += batch_index * m * c; 97 | 98 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 99 | const int stride = blockDim.y * blockDim.x; 100 | for (int i = index; i < c * n; i += stride) 101 | { 102 | const int l = i / n; 103 | const int j = i % n; 104 | float w1 = weight[j * 3 + 0]; 105 | float w2 = weight[j * 3 + 1]; 106 | float w3 = weight[j * 3 + 2]; 107 | int i1 = idx[j * 3 + 0]; 108 | int i2 = idx[j * 3 + 1]; 109 | int i3 = idx[j * 3 + 2]; 110 | atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); 111 | atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); 112 | atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); 113 | } 114 | } 115 | 116 | void nearestneighbor_cuda_launcher(int b, int n, int m, const float *unknown, const float *known, float *dist2, int *idx) 117 | { 118 | nearestneighbor_cuda_kernel<<>>(b, n, m, unknown, known, dist2, idx); 119 | } 120 | 121 | void interpolation_forward_cuda_launcher(int b, int c, int m, int n, const float *points, const int *idx, const float *weight, float *out) 122 | { 123 | interpolation_forward_cuda_kernel<<>>(b, c, m, n, points, idx, weight, out); 124 | } 125 | 126 | void interpolation_backward_cuda_launcher(int b, int n, int c, int m, const float *grad_out, const int *idx, const float *weight, float *grad_points) 127 | { 128 | interpolation_backward_cuda_kernel<<>>(b, n, c, m, grad_out, idx, weight, grad_points); 129 | } 130 | 131 | 132 | // input: unknown(b, n, 3) known(b, m, 3) 133 | // output: dist2(b, n, 3), idx(b, n, 3) 134 | __global__ void nearestneighbor_cuda_kernel_fast(int b, int n, int m, const float *__restrict__ unknown, const float *__restrict__ known, float *__restrict__ dist2, int *__restrict__ idx) { 135 | 136 | int bs_idx = blockIdx.y; 137 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 138 | if (bs_idx >= b || pt_idx >= n) return; 139 | 140 | unknown += bs_idx * n * 3 + pt_idx * 3; 141 | known += bs_idx * m * 3; 142 | dist2 += bs_idx * n * 3 + pt_idx * 3; 143 | idx += bs_idx * n * 3 + pt_idx * 3; 144 | 145 | float ux = unknown[0]; 146 | float uy = unknown[1]; 147 | float uz = unknown[2]; 148 | 149 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 150 | int besti1 = 0, besti2 = 0, besti3 = 0; 151 | for (int k = 0; k < m; ++k) { 152 | float x = known[k * 3 + 0]; 153 | float y = known[k * 3 + 1]; 154 | float z = known[k * 3 + 2]; 155 | float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 156 | if (d < best1) { 157 | best3 = best2; besti3 = besti2; 158 | best2 = best1; besti2 = besti1; 159 | best1 = d; besti1 = k; 160 | } 161 | else if (d < best2) { 162 | best3 = best2; besti3 = besti2; 163 | best2 = d; besti2 = k; 164 | } 165 | else if (d < best3) { 166 | best3 = d; besti3 = k; 167 | } 168 | } 169 | dist2[0] = best1; 170 | dist2[1] = best2; 171 | dist2[2] = best3; 172 | 173 | idx[0] = besti1; 174 | idx[1] = besti2; 175 | idx[2] = besti3; 176 | } 177 | 178 | 179 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 180 | // output: out(b, c, n) 181 | __global__ void interpolation_forward_cuda_kernel_fast(int b, int c, int m, int n, const float *__restrict__ points, const int *__restrict__ idx, const float *__restrict__ weight, float *__restrict__ out) { 182 | 183 | int bs_idx = blockIdx.z; 184 | int c_idx = blockIdx.y; 185 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 186 | 187 | if (bs_idx >= b || c_idx >= c || pt_idx >= n) return; 188 | 189 | weight += bs_idx * n * 3 + pt_idx * 3; 190 | points += bs_idx * c * m + c_idx * m; 191 | idx += bs_idx * n * 3 + pt_idx * 3; 192 | out += bs_idx * c * n + c_idx * n; 193 | 194 | out[pt_idx] = weight[0] * points[idx[0]] + weight[1] * points[idx[1]] + weight[2] * points[idx[2]]; 195 | } 196 | 197 | 198 | void nearestneighbor_cuda_launcher_fast(int b, int n, int m, const float *unknown, const float *known, float *dist2, int *idx) 199 | { 200 | cudaError_t err; 201 | 202 | dim3 blocks(DIVUP(n, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 203 | dim3 threads(THREADS_PER_BLOCK); 204 | 205 | nearestneighbor_cuda_kernel_fast<<>>(b, n, m, unknown, known, dist2, idx); 206 | 207 | err = cudaGetLastError(); 208 | if (cudaSuccess != err) { 209 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 210 | exit(-1); 211 | } 212 | } 213 | 214 | void interpolation_forward_cuda_launcher_fast(int b, int c, int m, int n, const float *points, const int *idx, const float *weight, float *out) { 215 | 216 | cudaError_t err; 217 | 218 | dim3 blocks(DIVUP(n, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) 219 | dim3 threads(THREADS_PER_BLOCK); 220 | interpolation_forward_cuda_kernel_fast<<>>(b, c, m, n, points, idx, weight, out); 221 | 222 | err = cudaGetLastError(); 223 | if (cudaSuccess != err) { 224 | fprintf(stderr, "CUDA kernel failed : %s\n", 225 | cudaGetErrorString(err)); 226 | exit(-1); 227 | } 228 | } 229 | -------------------------------------------------------------------------------- /libs/pointops/src/interpolation/interpolation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _INTERPOLATION_CUDA_KERNEL 2 | #define _INTERPOLATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void nearestneighbor_cuda(int b, int n, int m, at::Tensor unknown_tensor, at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor); 8 | void interpolation_forward_cuda(int b, int c, int m, int n, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor); 9 | void interpolation_backward_cuda(int b, int c, int n, int m, at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_points_tensor); 10 | 11 | void nearestneighbor_cuda_fast(int b, int n, int m, at::Tensor unknown_tensor, at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor); 12 | void interpolation_forward_cuda_fast(int b, int c, int m, int n, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor); 13 | 14 | #ifdef __cplusplus 15 | extern "C" { 16 | #endif 17 | 18 | void nearestneighbor_cuda_launcher(int b, int n, int m, const float *unknown, const float *known, float *dist2, int *idx); 19 | void interpolation_forward_cuda_launcher(int b, int c, int m, int n, const float *points, const int *idx, const float *weight, float *out); 20 | void interpolation_backward_cuda_launcher(int b, int c, int n, int m, const float *grad_out, const int *idx, const float *weight, float *grad_points); 21 | 22 | void nearestneighbor_cuda_launcher_fast(int b, int n, int m, const float *unknown, const float *known, float *dist2, int *idx); 23 | void interpolation_forward_cuda_launcher_fast(int b, int c, int m, int n, const float *points, const int *idx, const float *weight, float *out); 24 | 25 | #ifdef __cplusplus 26 | } 27 | #endif 28 | #endif 29 | -------------------------------------------------------------------------------- /libs/pointops/src/knnquery/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/libs/pointops/src/knnquery/__init__.py -------------------------------------------------------------------------------- /libs/pointops/src/knnquery/knnquery_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "knnquery_cuda_kernel.h" 7 | 8 | extern THCState *state; 9 | 10 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 11 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 12 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 13 | 14 | 15 | void knnquery_cuda(int b, int n, int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor) 16 | { 17 | CHECK_INPUT(new_xyz_tensor); 18 | CHECK_INPUT(xyz_tensor); 19 | 20 | const float *new_xyz = new_xyz_tensor.data(); 21 | const float *xyz = xyz_tensor.data(); 22 | int *idx = idx_tensor.data(); 23 | float *dist2 = dist2_tensor.data(); 24 | 25 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 26 | 27 | knnquery_cuda_launcher(b, n, m, nsample, xyz, new_xyz, idx, dist2, stream); 28 | } 29 | -------------------------------------------------------------------------------- /libs/pointops/src/knnquery/knnquery_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "knnquery_cuda_kernel.h" 3 | 4 | // input: xyz (b, n, 3) new_xyz (b, m, 3) 5 | // output: idx (b, m, nsample) dist2 (b, m, nsample) 6 | __global__ void knnquery_cuda_kernel(int b, int n, int m, int nsample, const float *__restrict__ xyz, const float *__restrict__ new_xyz, int *__restrict__ idx, float *__restrict__ dist2) { 7 | int bs_idx = blockIdx.y; 8 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 9 | if (bs_idx >= b || pt_idx >= m) return; 10 | 11 | new_xyz += bs_idx * m * 3 + pt_idx * 3; 12 | xyz += bs_idx * n * 3; 13 | idx += bs_idx * m * nsample + pt_idx * nsample; 14 | 15 | float new_x = new_xyz[0]; 16 | float new_y = new_xyz[1]; 17 | float new_z = new_xyz[2]; 18 | 19 | //double* best = new double[nsample]; 20 | //int* besti = new int[nsample]; 21 | double best[200]; 22 | int besti[200]; 23 | for(int i = 0; i < nsample; i++){ 24 | best[i] = 1e40; 25 | besti[i] = 0; 26 | } 27 | for(int k = 0; k < n; k++){ 28 | float x = xyz[k * 3 + 0]; 29 | float y = xyz[k * 3 + 1]; 30 | float z = xyz[k * 3 + 2]; 31 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 32 | for(int j = 0; j < nsample; j++){ 33 | if(d2 < best[j]){ 34 | for(int i = nsample - 1; i > j; i--){ 35 | best[i] = best[i - 1]; 36 | besti[i] = besti[i - 1]; 37 | } 38 | best[j] = d2; 39 | besti[j] = k; 40 | break; 41 | } 42 | } 43 | } 44 | for(int i = 0; i < nsample; i++){ 45 | idx[i] = besti[i]; 46 | dist2[i] = best[i]; 47 | } 48 | //delete []best; 49 | //delete []besti; 50 | } 51 | 52 | 53 | void knnquery_cuda_launcher(int b, int n, int m, int nsample, const float *xyz, const float *new_xyz, int *idx, float *dist2, cudaStream_t stream) { 54 | // param new_xyz: (B, m, 3) 55 | // param xyz: (B, n, 3) 56 | // param idx: (B, m, nsample) 57 | 58 | cudaError_t err; 59 | 60 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 61 | dim3 threads(THREADS_PER_BLOCK); 62 | 63 | knnquery_cuda_kernel<<>>(b, n, m, nsample, xyz, new_xyz, idx, dist2); 64 | // cudaDeviceSynchronize(); // for using printf in kernel function 65 | 66 | err = cudaGetLastError(); 67 | if (cudaSuccess != err) { 68 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 69 | exit(-1); 70 | } 71 | } -------------------------------------------------------------------------------- /libs/pointops/src/knnquery/knnquery_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _KNNQUERY_CUDA_KERNEL 2 | #define _KNNQUERY_CUDA_KERNEL 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | void knnquery_cuda(int b, int n, int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void knnquery_cuda_launcher(int b, int n, int m, int nsample, const float *xyz, const float *new_xyz, int *idx, float *dist2, cudaStream_t stream); 15 | 16 | #ifdef __cplusplus 17 | } 18 | #endif 19 | 20 | #endif -------------------------------------------------------------------------------- /libs/pointops/src/labelstat/labelstat_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "labelstat_cuda_kernel.h" 7 | 8 | extern THCState *state; 9 | 10 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 11 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 12 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 13 | 14 | void labelstat_idx_cuda_fast(int b, int n, int m, int nsample, int nclass, 15 | at::Tensor label_stat_tensor, at::Tensor idx_tensor, at::Tensor new_label_stat_tensor) 16 | { 17 | CHECK_INPUT(label_stat_tensor); 18 | CHECK_INPUT(idx_tensor); 19 | 20 | const int *label_stat = label_stat_tensor.data(); 21 | const int *idx = idx_tensor.data(); 22 | int *new_label_stat = new_label_stat_tensor.data(); 23 | 24 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 25 | 26 | labelstat_idx_cuda_launcher_fast(b, n, m, nsample, nclass, label_stat, idx, new_label_stat, stream); 27 | } 28 | 29 | void labelstat_ballrange_cuda_fast(int b, int n, int m, float radius, int nclass, 30 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor label_stat_tensor, at::Tensor new_label_stat_tensor) 31 | { 32 | CHECK_INPUT(new_xyz_tensor); 33 | CHECK_INPUT(xyz_tensor); 34 | CHECK_INPUT(label_stat_tensor); 35 | 36 | const float *new_xyz = new_xyz_tensor.data(); 37 | const float *xyz = xyz_tensor.data(); 38 | const int *label_stat = label_stat_tensor.data(); 39 | int *new_label_stat = new_label_stat_tensor.data(); 40 | 41 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 42 | 43 | labelstat_ballrange_cuda_launcher_fast(b, n, m, radius, nclass, new_xyz, xyz, label_stat, new_label_stat, stream); 44 | } 45 | 46 | void labelstat_and_ballquery_cuda_fast(int b, int n, int m, float radius, int nsample, int nclass, 47 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor label_stat_tensor, at::Tensor idx_tensor, at::Tensor new_label_stat_tensor) 48 | { 49 | CHECK_INPUT(new_xyz_tensor); 50 | CHECK_INPUT(xyz_tensor); 51 | CHECK_INPUT(label_stat_tensor); 52 | CHECK_INPUT(idx_tensor); 53 | 54 | const float *new_xyz = new_xyz_tensor.data(); 55 | const float *xyz = xyz_tensor.data(); 56 | const int *label_stat = label_stat_tensor.data(); 57 | int *idx = idx_tensor.data(); 58 | int *new_label_stat = new_label_stat_tensor.data(); 59 | 60 | cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); 61 | 62 | labelstat_and_ballquery_cuda_launcher_fast(b, n, m, radius, nsample, nclass, new_xyz, xyz, label_stat, idx, new_label_stat, stream); 63 | } 64 | -------------------------------------------------------------------------------- /libs/pointops/src/labelstat/labelstat_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "labelstat_cuda_kernel.h" 3 | 4 | // input: new_xyz(b, m, 3) xyz(b, n, 3) label_stat(b, n, nclass) 5 | // output: idx(b, m, nsample) new_label_stat(b, m, nclass) 6 | __global__ void labelstat_and_ballquery_cuda_kernel_fast(int b, int n, int m, float radius, int nsample, int nclass, 7 | const float *new_xyz, const float *xyz, const int *label_stat, int *idx, int *new_label_stat) { 8 | int bs_idx = blockIdx.y; 9 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 10 | if (bs_idx >= b || pt_idx >= m) return; 11 | 12 | new_xyz += bs_idx * m * 3 + pt_idx * 3; 13 | xyz += bs_idx * n * 3; 14 | idx += bs_idx * m * nsample + pt_idx * nsample; 15 | label_stat += bs_idx * n * nclass; 16 | new_label_stat += bs_idx * m * nclass + pt_idx * nclass; 17 | 18 | for(int i = 0; i < nclass; i++){ 19 | new_label_stat[i] = 0; 20 | } 21 | 22 | float radius2 = radius * radius; 23 | float new_x = new_xyz[0]; 24 | float new_y = new_xyz[1]; 25 | float new_z = new_xyz[2]; 26 | 27 | int cnt = 0; 28 | for (int k = 0; k < n; ++k) { 29 | float x = xyz[k * 3 + 0]; 30 | float y = xyz[k * 3 + 1]; 31 | float z = xyz[k * 3 + 2]; 32 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 33 | if (d2 < radius2){ 34 | for(int i = 0; i < nclass; i++){ 35 | new_label_stat[i] += label_stat[k * nclass + i]; 36 | } 37 | if (cnt == 0){ 38 | for (int l = 0; l < nsample; ++l) { 39 | idx[l] = k; 40 | } 41 | } 42 | idx[cnt] = k; 43 | ++cnt; 44 | if (cnt >= nsample){ 45 | break; 46 | } 47 | } 48 | } 49 | } 50 | 51 | void labelstat_and_ballquery_cuda_launcher_fast(int b, int n, int m, float radius, int nsample, int nclass, 52 | const float *new_xyz, const float *xyz, const int *label_stat, int *idx, int *new_label_stat, cudaStream_t stream) { 53 | // param new_xyz: (B, m, 3) 54 | // param xyz: (B, n, 3) 55 | // param idx: (B, m, nsample) 56 | 57 | cudaError_t err; 58 | 59 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 60 | dim3 threads(THREADS_PER_BLOCK); 61 | 62 | labelstat_and_ballquery_cuda_kernel_fast<<>>(b, n, m, radius, nsample, nclass, new_xyz, xyz, label_stat, idx, new_label_stat); 63 | // cudaDeviceSynchronize(); // for using printf in kernel function 64 | 65 | err = cudaGetLastError(); 66 | if (cudaSuccess != err) { 67 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 68 | exit(-1); 69 | } 70 | } 71 | 72 | // input: new_xyz(b, m, 3) xyz(b, n, 3) label_stat(b, n, nclass) 73 | // output: new_label_stat(b, m, nclass) 74 | __global__ void labelstat_ballrange_cuda_kernel_fast(int b, int n, int m, float radius, int nclass, 75 | const float *new_xyz, const float *xyz, const int *label_stat, int *new_label_stat) { 76 | int bs_idx = blockIdx.y; 77 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 78 | if (bs_idx >= b || pt_idx >= m) return; 79 | 80 | new_xyz += bs_idx * m * 3 + pt_idx * 3; 81 | xyz += bs_idx * n * 3; 82 | label_stat += bs_idx * n * nclass; 83 | new_label_stat += bs_idx * m * nclass + pt_idx * nclass; 84 | 85 | for(int i = 0; i < nclass; i++){ 86 | new_label_stat[i] = 0; 87 | } 88 | 89 | float radius2 = radius * radius; 90 | float new_x = new_xyz[0]; 91 | float new_y = new_xyz[1]; 92 | float new_z = new_xyz[2]; 93 | 94 | for (int k = 0; k < n; ++k) { 95 | float x = xyz[k * 3 + 0]; 96 | float y = xyz[k * 3 + 1]; 97 | float z = xyz[k * 3 + 2]; 98 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 99 | if (d2 < radius2){ 100 | for(int i = 0; i < nclass; i++){ 101 | new_label_stat[i] += label_stat[k * nclass + i]; 102 | } 103 | } 104 | } 105 | } 106 | 107 | 108 | void labelstat_ballrange_cuda_launcher_fast(int b, int n, int m, float radius, int nclass, 109 | const float *new_xyz, const float *xyz, const int *label_stat, int *new_label_stat, cudaStream_t stream) { 110 | // param new_xyz: (B, m, 3) 111 | // param xyz: (B, n, 3) 112 | // param idx: (B, m, nsample) 113 | 114 | cudaError_t err; 115 | 116 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 117 | dim3 threads(THREADS_PER_BLOCK); 118 | 119 | labelstat_ballrange_cuda_kernel_fast<<>>(b, n, m, radius, nclass, new_xyz, xyz, label_stat, new_label_stat); 120 | // cudaDeviceSynchronize(); // for using printf in kernel function 121 | 122 | err = cudaGetLastError(); 123 | if (cudaSuccess != err) { 124 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 125 | exit(-1); 126 | } 127 | } 128 | 129 | // input: idx(b, m, nsample) label_stat(b, n, nclass) 130 | // output: new_label_stat(b, m, nclass) 131 | __global__ void labelstat_idx_cuda_kernel_fast(int b, int n, int m, int nsample, int nclass, 132 | const int *label_stat, const int *idx, int *new_label_stat) { 133 | int bs_idx = blockIdx.y; 134 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 135 | if (bs_idx >= b || pt_idx >= m) return; 136 | 137 | idx += bs_idx * m * nsample + pt_idx * nsample; 138 | label_stat += bs_idx * n * nclass; 139 | new_label_stat += bs_idx * m * nclass + pt_idx * nclass; 140 | 141 | for(int i = 0; i < nclass; i++){ 142 | new_label_stat[i] = 0; 143 | } 144 | 145 | for(int k = 0; k < nsample; k++){ 146 | const int *label_stat_k = label_stat + idx[k] * nclass; 147 | for(int i = 0; i < nclass; i++){ 148 | new_label_stat[i] += label_stat_k[i]; 149 | } 150 | } 151 | } 152 | 153 | 154 | void labelstat_idx_cuda_launcher_fast(int b, int n, int m, int nsample, int nclass, 155 | const int *label_stat, const int *idx, int *new_label_stat, cudaStream_t stream) { 156 | // param new_xyz: (B, m, 3) 157 | // param xyz: (B, n, 3) 158 | // param idx: (B, m, nsample) 159 | 160 | cudaError_t err; 161 | 162 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 163 | dim3 threads(THREADS_PER_BLOCK); 164 | 165 | labelstat_idx_cuda_kernel_fast<<>>(b, n, m, nsample, nclass, label_stat, idx, new_label_stat); 166 | // cudaDeviceSynchronize(); // for using printf in kernel function 167 | 168 | err = cudaGetLastError(); 169 | if (cudaSuccess != err) { 170 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 171 | exit(-1); 172 | } 173 | } -------------------------------------------------------------------------------- /libs/pointops/src/labelstat/labelstat_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _LABELSTAT_CUDA_KERNEL 2 | #define _LABELSTAT_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void labelstat_and_ballquery_cuda_fast(int b, int n, int m, float radius, int nsample, int nclass, 8 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor label_stat_tensor, at::Tensor idx_tensor, at::Tensor new_label_stat_tensor); 9 | 10 | void labelstat_ballrange_cuda_fast(int b, int n, int m, float radius, int nclass, 11 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor label_stat_tensor, at::Tensor new_label_stat_tensor); 12 | 13 | void labelstat_idx_cuda_fast(int b, int n, int m, int nsample, int nclass, 14 | at::Tensor label_stat_tensor, at::Tensor idx_tensor, at::Tensor new_label_stat_tensor); 15 | 16 | #ifdef __cplusplus 17 | extern "C" { 18 | #endif 19 | 20 | void labelstat_and_ballquery_cuda_launcher_fast(int b, int n, int m, float radius, int nsample, int nclass, \ 21 | const float *new_xyz, const float *xyz, const int *label_stat, int *idx, int *new_label_stat, cudaStream_t stream); 22 | 23 | void labelstat_ballrange_cuda_launcher_fast(int b, int n, int m, float radius, int nclass, \ 24 | const float *new_xyz, const float *xyz, const int *label_stat, int *new_label_stat, cudaStream_t stream); 25 | 26 | void labelstat_idx_cuda_launcher_fast(int b, int n, int m, int nsample, int nclass, \ 27 | const int *label_stat, const int *idx, int *new_label_stat, cudaStream_t stream); 28 | 29 | #ifdef __cplusplus 30 | } 31 | #endif 32 | 33 | #endif 34 | -------------------------------------------------------------------------------- /libs/pointops/src/pointops_api.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "ballquery/ballquery_cuda_kernel.h" 5 | #include "grouping/grouping_cuda_kernel.h" 6 | #include "grouping_int/grouping_int_cuda_kernel.h" 7 | #include "sampling/sampling_cuda_kernel.h" 8 | #include "interpolation/interpolation_cuda_kernel.h" 9 | #include "knnquery/knnquery_cuda_kernel.h" 10 | 11 | #include "labelstat/labelstat_cuda_kernel.h" 12 | #include "featuredistribute/featuredistribute_cuda_kernel.h" 13 | 14 | 15 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 16 | m.def("ballquery_cuda", &ballquery_cuda_fast, "ballquery_cuda_fast"); // name in python, cpp function address, docs 17 | 18 | m.def("knnquery_cuda", &knnquery_cuda, "knnquery_cuda"); 19 | 20 | m.def("grouping_forward_cuda", &grouping_forward_cuda_fast, "grouping_forward_cuda_fast"); 21 | m.def("grouping_backward_cuda", &grouping_backward_cuda, "grouping_backward_cuda"); 22 | 23 | m.def("grouping_int_forward_cuda", &grouping_int_forward_cuda_fast, "grouping_int_forward_cuda_fast"); 24 | 25 | m.def("gathering_forward_cuda", &gathering_forward_cuda, "gathering_forward_cuda"); 26 | m.def("gathering_backward_cuda", &gathering_backward_cuda, "gathering_backward_cuda"); 27 | m.def("furthestsampling_cuda", &furthestsampling_cuda, "furthestsampling_cuda"); 28 | 29 | m.def("nearestneighbor_cuda", &nearestneighbor_cuda_fast, "nearestneighbor_cuda_fast"); 30 | m.def("interpolation_forward_cuda", &interpolation_forward_cuda_fast, "interpolation_forward_cuda_fast"); 31 | m.def("interpolation_backward_cuda", &interpolation_backward_cuda, "interpolation_backward_cuda"); 32 | 33 | m.def("labelstat_idx_cuda", &labelstat_idx_cuda_fast, "labelstat_idx_cuda_fast"); 34 | m.def("labelstat_ballrange_cuda", &labelstat_ballrange_cuda_fast, "labelstat_ballrange_cuda_fast"); 35 | m.def("labelstat_and_ballquery_cuda", &labelstat_and_ballquery_cuda_fast, "labelstat_and_ballquery_cuda_fast"); 36 | 37 | m.def("featuredistribute_cuda", &featuredistribute_cuda, "featuredistribute_cuda"); 38 | m.def("featuregather_forward_cuda", &featuregather_forward_cuda, "featuregather_forward_cuda"); 39 | m.def("featuregather_backward_cuda", &featuregather_backward_cuda, "featuregather_backward_cuda"); 40 | } 41 | -------------------------------------------------------------------------------- /libs/pointops/src/sampling/sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "sampling_cuda_kernel.h" 6 | 7 | extern THCState *state; 8 | 9 | void gathering_forward_cuda(int b, int c, int n, int m, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) 10 | { 11 | const float *points = points_tensor.data(); 12 | const int *idx = idx_tensor.data(); 13 | float *out = out_tensor.data(); 14 | gathering_forward_cuda_launcher(b, c, n, m, points, idx, out); 15 | } 16 | 17 | void gathering_backward_cuda(int b, int c, int n, int m, at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) 18 | { 19 | 20 | const float *grad_out = grad_out_tensor.data(); 21 | const int *idx = idx_tensor.data(); 22 | float *grad_points = grad_points_tensor.data(); 23 | gathering_backward_cuda_launcher(b, c, n, m, grad_out, idx, grad_points); 24 | } 25 | 26 | void furthestsampling_cuda(int b, int n, int m, at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) 27 | { 28 | const float *points = points_tensor.data(); 29 | float *temp = temp_tensor.data(); 30 | int *idx = idx_tensor.data(); 31 | furthestsampling_cuda_launcher(b, n, m, points, temp, idx); 32 | } 33 | -------------------------------------------------------------------------------- /libs/pointops/src/sampling/sampling_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "sampling_cuda_kernel.h" 3 | 4 | // input: points(b, c, n) idx(b, m) 5 | // output: out(b, c, m) 6 | __global__ void gathering_forward_cuda_kernel(int b, int c, int n, int m, const float *points, const int *idx, float *out) 7 | { 8 | for (int i = blockIdx.x; i < b; i += gridDim.x) 9 | { 10 | for (int l = blockIdx.y; l < c; l += gridDim.y) 11 | { 12 | for (int j = threadIdx.x; j < m; j += blockDim.x) 13 | { 14 | int a = idx[i * m + j]; 15 | out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; 16 | } 17 | } 18 | } 19 | } 20 | 21 | // input: grad_out(b, c, m) idx(b, m) 22 | // output: grad_points(b, c, n) 23 | __global__ void gathering_backward_cuda_kernel(int b, int c, int n, int m, const float *grad_out, const int *idx, float *grad_points) 24 | { 25 | for (int i = blockIdx.x; i < b; i += gridDim.x) 26 | { 27 | for (int l = blockIdx.y; l < c; l += gridDim.y) 28 | { 29 | for (int j = threadIdx.x; j < m; j += blockDim.x) 30 | { 31 | int a = idx[i * m + j]; 32 | atomicAdd(grad_points + (i * c + l) * n + a, grad_out[(i * c + l) * m + j]); 33 | } 34 | } 35 | } 36 | } 37 | 38 | void gathering_forward_cuda_launcher(int b, int c, int n, int m, const float *points, const int *idx, float *out) 39 | { 40 | gathering_forward_cuda_kernel<<>>(b, c, n, m, points, idx, out); 41 | } 42 | 43 | void gathering_backward_cuda_launcher(int b, int c, int n, int m, const float *grad_out, const int *idx, float *grad_points) 44 | { 45 | gathering_backward_cuda_kernel<<>>(b, c, n, m, grad_out, idx, grad_points); 46 | } 47 | 48 | __device__ void __update(float *dists, int *dists_i, 49 | int idx1, int idx2) { 50 | const float v1 = dists[idx1], v2 = dists[idx2]; 51 | const int i1 = dists_i[idx1], i2 = dists_i[idx2]; 52 | dists[idx1] = max(v1, v2); 53 | dists_i[idx1] = v2 > v1 ? i2 : i1; 54 | } 55 | 56 | // Input dataset: (b, n, 3), tmp: (b, n) 57 | // Ouput idxs (b, m) 58 | template 59 | __global__ void furthestsampling_cuda_kernel(int b, int n, int m, const float *dataset, float *temp, int *idxs) 60 | { 61 | if (m <= 0) 62 | return; 63 | __shared__ float dists[block_size]; 64 | __shared__ int dists_i[block_size]; 65 | 66 | int batch_index = blockIdx.x; 67 | dataset += batch_index * n * 3; 68 | temp += batch_index * n; 69 | idxs += batch_index * m; 70 | int tid = threadIdx.x; 71 | const int stride = block_size; 72 | int old = 0; 73 | if (threadIdx.x == 0) 74 | idxs[0] = old; 75 | 76 | __syncthreads(); 77 | for (int j = 1; j < m; j++) 78 | { 79 | int besti = 0; 80 | float best = -1; 81 | float x1 = dataset[old * 3 + 0]; 82 | float y1 = dataset[old * 3 + 1]; 83 | float z1 = dataset[old * 3 + 2]; 84 | for (int k = tid; k < n; k += stride) 85 | { 86 | float x2, y2, z2; 87 | x2 = dataset[k * 3 + 0]; 88 | y2 = dataset[k * 3 + 1]; 89 | z2 = dataset[k * 3 + 2]; 90 | //float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); 91 | //if (mag <= 1e-3) 92 | // continue; 93 | float d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); 94 | float d2 = min(d, temp[k]); 95 | temp[k] = d2; 96 | besti = d2 > best ? k : besti; 97 | best = d2 > best ? d2 : best; 98 | } 99 | dists[tid] = best; 100 | dists_i[tid] = besti; 101 | __syncthreads(); 102 | 103 | if (block_size >= 1024) { 104 | if (tid < 512) { 105 | __update(dists, dists_i, tid, tid + 512); 106 | } 107 | __syncthreads(); 108 | } 109 | if (block_size >= 512) { 110 | if (tid < 256) { 111 | __update(dists, dists_i, tid, tid + 256); 112 | } 113 | __syncthreads(); 114 | } 115 | if (block_size >= 256) { 116 | if (tid < 128) { 117 | __update(dists, dists_i, tid, tid + 128); 118 | } 119 | __syncthreads(); 120 | } 121 | if (block_size >= 128) { 122 | if (tid < 64) { 123 | __update(dists, dists_i, tid, tid + 64); 124 | } 125 | __syncthreads(); 126 | } 127 | if (block_size >= 64) { 128 | if (tid < 32) { 129 | __update(dists, dists_i, tid, tid + 32); 130 | } 131 | __syncthreads(); 132 | } 133 | if (block_size >= 32) { 134 | if (tid < 16) { 135 | __update(dists, dists_i, tid, tid + 16); 136 | } 137 | __syncthreads(); 138 | } 139 | if (block_size >= 16) { 140 | if (tid < 8) { 141 | __update(dists, dists_i, tid, tid + 8); 142 | } 143 | __syncthreads(); 144 | } 145 | if (block_size >= 8) { 146 | if (tid < 4) { 147 | __update(dists, dists_i, tid, tid + 4); 148 | } 149 | __syncthreads(); 150 | } 151 | if (block_size >= 4) { 152 | if (tid < 2) { 153 | __update(dists, dists_i, tid, tid + 2); 154 | } 155 | __syncthreads(); 156 | } 157 | if (block_size >= 2) { 158 | if (tid < 1) { 159 | __update(dists, dists_i, tid, tid + 1); 160 | } 161 | __syncthreads(); 162 | } 163 | 164 | old = dists_i[0]; 165 | if (tid == 0) 166 | idxs[j] = old; 167 | } 168 | } 169 | 170 | void furthestsampling_cuda_launcher(int b, int n, int m, const float *dataset, float *temp, int *idxs) 171 | { 172 | unsigned int n_threads = opt_n_threads(n); 173 | switch (n_threads) { 174 | case 1024: 175 | furthestsampling_cuda_kernel<1024><<>>(b, n, m, dataset, temp, idxs); 176 | break; 177 | case 512: 178 | furthestsampling_cuda_kernel<512><<>>(b, n, m, dataset, temp, idxs); 179 | break; 180 | case 256: 181 | furthestsampling_cuda_kernel<256><<>>(b, n, m, dataset, temp, idxs); 182 | break; 183 | case 128: 184 | furthestsampling_cuda_kernel<128><<>>(b, n, m, dataset, temp, idxs); 185 | break; 186 | case 64: 187 | furthestsampling_cuda_kernel<64><<>>(b, n, m, dataset, temp, idxs); 188 | break; 189 | case 32: 190 | furthestsampling_cuda_kernel<32><<>>(b, n, m, dataset, temp, idxs); 191 | break; 192 | case 16: 193 | furthestsampling_cuda_kernel<16><<>>(b, n, m, dataset, temp, idxs); 194 | break; 195 | case 8: 196 | furthestsampling_cuda_kernel<8><<>>(b, n, m, dataset, temp, idxs); 197 | break; 198 | case 4: 199 | furthestsampling_cuda_kernel<4><<>>(b, n, m, dataset, temp, idxs); 200 | break; 201 | case 2: 202 | furthestsampling_cuda_kernel<2><<>>(b, n, m, dataset, temp, idxs); 203 | break; 204 | case 1: 205 | furthestsampling_cuda_kernel<1><<>>(b, n, m, dataset, temp, idxs); 206 | break; 207 | default: 208 | furthestsampling_cuda_kernel<512><<>>(b, n, m, dataset, temp, idxs); 209 | } 210 | } 211 | -------------------------------------------------------------------------------- /libs/pointops/src/sampling/sampling_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_CUDA_KERNEL 2 | #define _SAMPLING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void gathering_forward_cuda(int b, int c, int n, int m, at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor); 8 | void gathering_backward_cuda(int b, int c, int n, int m, at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor); 9 | void furthestsampling_cuda(int b, int n, int m, at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor); 10 | 11 | #ifdef __cplusplus 12 | extern "C" { 13 | #endif 14 | 15 | void gathering_forward_cuda_launcher(int b, int c, int n, int m, const float *points, const int *idx, float *out); 16 | void gathering_backward_cuda_launcher(int b, int c, int n, int m, const float *grad_out, const int *idx, float *grad_points); 17 | void furthestsampling_cuda_launcher(int b, int n, int m, const float *dataset, float *temp, int *idxs); 18 | 19 | #ifdef __cplusplus 20 | } 21 | #endif 22 | #endif 23 | -------------------------------------------------------------------------------- /loss/__pycache__/cosface_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/loss/__pycache__/cosface_loss.cpython-36.pyc -------------------------------------------------------------------------------- /loss/__pycache__/cosface_loss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/loss/__pycache__/cosface_loss.cpython-38.pyc -------------------------------------------------------------------------------- /loss/__pycache__/incremental_loss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/loss/__pycache__/incremental_loss.cpython-38.pyc -------------------------------------------------------------------------------- /loss/__pycache__/triplet_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/loss/__pycache__/triplet_loss.cpython-36.pyc -------------------------------------------------------------------------------- /loss/__pycache__/triplet_loss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/loss/__pycache__/triplet_loss.cpython-38.pyc -------------------------------------------------------------------------------- /loss/incremental_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from tqdm import tqdm 5 | 6 | class NoIncLoss: 7 | def __init__(self): 8 | pass 9 | 10 | def adjust_weight(self, epoch): 11 | pass 12 | 13 | def __call__(self, *args, **kwargs): 14 | return torch.tensor(0, dtype = float, device = 'cuda') 15 | 16 | class LwF: 17 | def __init__(self): 18 | self.weight = 1000 19 | self.temperature = 2 20 | 21 | def adjust_weight(self, epoch): 22 | pass 23 | 24 | def __call__(self, old_rep, new_rep): 25 | log_p = torch.log_softmax(new_rep / self.temperature, dim=1) 26 | q = torch.softmax(old_rep / self.temperature, dim=1) 27 | res = torch.nn.functional.kl_div(log_p, q, reduction="batchmean") 28 | loss_incremental = self.weight * res 29 | return loss_incremental 30 | 31 | -------------------------------------------------------------------------------- /loss/triplet_loss.py: -------------------------------------------------------------------------------- 1 | # Warsaw University of Technology 2 | import torch 3 | from pytorch_metric_learning import losses, reducers 4 | from pytorch_metric_learning.distances import LpDistance 5 | 6 | 7 | 8 | class HardTripletMinerWithMasks: 9 | # Hard triplet miner 10 | def __init__(self, distance): 11 | self.distance = distance 12 | # Stats 13 | self.max_pos_pair_dist = None 14 | self.max_neg_pair_dist = None 15 | self.mean_pos_pair_dist = None 16 | self.mean_neg_pair_dist = None 17 | self.min_pos_pair_dist = None 18 | self.min_neg_pair_dist = None 19 | 20 | def __call__(self, embeddings, positives_mask, negatives_mask): 21 | assert embeddings.dim() == 2 22 | d_embeddings = embeddings.detach() 23 | with torch.no_grad(): 24 | hard_triplets = self.mine(d_embeddings, positives_mask, negatives_mask) 25 | return hard_triplets 26 | 27 | def mine(self, embeddings, positives_mask, negatives_mask): 28 | # Based on pytorch-metric-learning implementation 29 | dist_mat = self.distance(embeddings) 30 | (hardest_positive_dist, hardest_positive_indices), a1p_keep = get_max_per_row(dist_mat, positives_mask) 31 | (hardest_negative_dist, hardest_negative_indices), a2n_keep = get_min_per_row(dist_mat, negatives_mask) 32 | a_keep_idx = torch.where(a1p_keep & a2n_keep) 33 | a = torch.arange(dist_mat.size(0)).to(hardest_positive_indices.device)[a_keep_idx] 34 | p = hardest_positive_indices[a_keep_idx] 35 | n = hardest_negative_indices[a_keep_idx] 36 | self.max_pos_pair_dist = torch.max(hardest_positive_dist[a_keep_idx]).item() 37 | self.max_neg_pair_dist = torch.max(hardest_negative_dist[a_keep_idx]).item() 38 | self.mean_pos_pair_dist = torch.mean(hardest_positive_dist[a_keep_idx]).item() 39 | self.mean_neg_pair_dist = torch.mean(hardest_negative_dist[a_keep_idx]).item() 40 | self.min_pos_pair_dist = torch.min(hardest_positive_dist[a_keep_idx]).item() 41 | self.min_neg_pair_dist = torch.min(hardest_negative_dist[a_keep_idx]).item() 42 | return a, p, n 43 | 44 | 45 | def get_max_per_row(mat, mask): 46 | non_zero_rows = torch.any(mask, dim=1) 47 | mat_masked = mat.clone() 48 | mat_masked[~mask] = 0 49 | return torch.max(mat_masked, dim=1), non_zero_rows 50 | 51 | 52 | def get_min_per_row(mat, mask): 53 | non_inf_rows = torch.any(mask, dim=1) 54 | mat_masked = mat.clone() 55 | mat_masked[~mask] = float('inf') 56 | return torch.min(mat_masked, dim=1), non_inf_rows 57 | 58 | 59 | class BatchHardTripletLossWithMasks: 60 | def __init__(self, margin:float): 61 | self.margin = margin 62 | self.distance = LpDistance(normalize_embeddings=False, collect_stats=True) 63 | # We use triplet loss with Euclidean distance 64 | self.miner_fn = HardTripletMinerWithMasks(distance=self.distance) 65 | reducer_fn = reducers.AvgNonZeroReducer(collect_stats=True) 66 | self.loss_fn = losses.TripletMarginLoss(margin=self.margin, swap=True, distance=self.distance, 67 | reducer=reducer_fn, collect_stats=True) 68 | 69 | def __call__(self, embeddings, positives_mask, negatives_mask): 70 | hard_triplets = self.miner_fn(embeddings, positives_mask, negatives_mask) 71 | dummy_labels = torch.arange(embeddings.shape[0]).to(embeddings.device) 72 | loss = self.loss_fn(embeddings, dummy_labels, hard_triplets) 73 | 74 | stats = {'loss': loss.item(), 'avg_embedding_norm': self.loss_fn.distance.final_avg_query_norm, 75 | 'num_non_zero_triplets': self.loss_fn.reducer.triplets_past_filter, 76 | 'num_triplets': len(hard_triplets[0]), 77 | 'mean_pos_pair_dist': self.miner_fn.mean_pos_pair_dist, 78 | 'mean_neg_pair_dist': self.miner_fn.mean_neg_pair_dist, 79 | 'max_pos_pair_dist': self.miner_fn.max_pos_pair_dist, 80 | 'max_neg_pair_dist': self.miner_fn.max_neg_pair_dist, 81 | 'min_pos_pair_dist': self.miner_fn.min_pos_pair_dist, 82 | 'min_neg_pair_dist': self.miner_fn.min_neg_pair_dist 83 | } 84 | return loss, stats 85 | 86 | 87 | class BatchHardContrastiveLossWithMasks: 88 | def __init__(self, pos_margin: float, neg_margin: float): 89 | self.pos_margin = pos_margin 90 | self.neg_margin = neg_margin 91 | self.distance = LpDistance(normalize_embeddings=False, collect_stats=True) 92 | self.miner_fn = HardTripletMinerWithMasks(distance=self.distance) 93 | # We use contrastive loss with squared Euclidean distance 94 | reducer_fn = reducers.AvgNonZeroReducer(collect_stats=True) 95 | self.loss_fn = losses.ContrastiveLoss(pos_margin=self.pos_margin, neg_margin=self.neg_margin, 96 | distance=self.distance, reducer=reducer_fn, collect_stats=True) 97 | 98 | def __call__(self, embeddings, positives_mask, negatives_mask): 99 | hard_triplets = self.miner_fn(embeddings, positives_mask, negatives_mask) 100 | dummy_labels = torch.arange(embeddings.shape[0]).to(embeddings.device) 101 | loss = self.loss_fn(embeddings, dummy_labels, hard_triplets) 102 | stats = {'loss': loss.item(), 'avg_embedding_norm': self.loss_fn.distance.final_avg_query_norm, 103 | 'pos_pairs_above_threshold': self.loss_fn.reducer.reducers['pos_loss'].pos_pairs_above_threshold, 104 | 'neg_pairs_above_threshold': self.loss_fn.reducer.reducers['neg_loss'].neg_pairs_above_threshold, 105 | 'pos_loss': self.loss_fn.reducer.reducers['pos_loss'].pos_loss.item(), 106 | 'neg_loss': self.loss_fn.reducer.reducers['neg_loss'].neg_loss.item(), 107 | 'num_pairs': 2*len(hard_triplets[0]), 108 | 'mean_pos_pair_dist': self.miner_fn.mean_pos_pair_dist, 109 | 'mean_neg_pair_dist': self.miner_fn.mean_neg_pair_dist, 110 | 'max_pos_pair_dist': self.miner_fn.max_pos_pair_dist, 111 | 'max_neg_pair_dist': self.miner_fn.max_neg_pair_dist, 112 | 'min_pos_pair_dist': self.miner_fn.min_pos_pair_dist, 113 | 'min_neg_pair_dist': self.miner_fn.min_neg_pair_dist 114 | } 115 | 116 | return loss, stats 117 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/__pycache__/mink_params.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/__pycache__/mink_params.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/__pycache__/minkfpn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/__pycache__/minkfpn.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/__pycache__/minkloc.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/__pycache__/minkloc.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/__pycache__/resnet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/__pycache__/resnet.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/__pycache__/eca_block.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/layers/__pycache__/eca_block.cpython-36.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/__pycache__/eca_block.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/layers/__pycache__/eca_block.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/__pycache__/netvlad.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/layers/__pycache__/netvlad.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/__pycache__/pooling.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/layers/__pycache__/pooling.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/__pycache__/pooling_wrapper.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/layers/__pycache__/pooling_wrapper.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/eca_block.py: -------------------------------------------------------------------------------- 1 | # Implementation of Efficient Channel Attention ECA block 2 | 3 | import numpy as np 4 | import torch.nn as nn 5 | 6 | import MinkowskiEngine as ME 7 | 8 | from MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck 9 | 10 | 11 | class ECALayer(nn.Module): 12 | def __init__(self, channels, gamma=2, b=1): 13 | super().__init__() 14 | t = int(abs((np.log2(channels) + b) / gamma)) 15 | k_size = t if t % 2 else t + 1 16 | self.avg_pool = ME.MinkowskiGlobalPooling() 17 | self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False) 18 | self.sigmoid = nn.Sigmoid() 19 | self.broadcast_mul = ME.MinkowskiBroadcastMultiplication() 20 | 21 | def forward(self, x: ME.SparseTensor): 22 | # feature descriptor on the global spatial information 23 | y_sparse = self.avg_pool(x) 24 | 25 | # Apply 1D convolution along the channel dimension 26 | y = self.conv(y_sparse.F.unsqueeze(1)).squeeze(1) 27 | # y is (batch_size, channels) tensor 28 | 29 | y = self.sigmoid(y) 30 | # y is (batch_size, channels) tensor 31 | 32 | y_sparse = ME.SparseTensor(y, coordinate_manager=y_sparse.coordinate_manager, 33 | coordinate_map_key=y_sparse.coordinate_map_key) 34 | # y must be features reduced to the origin 35 | return self.broadcast_mul(x, y_sparse) 36 | 37 | 38 | class ECABasicBlock(BasicBlock): 39 | def __init__(self, 40 | inplanes, 41 | planes, 42 | stride=1, 43 | dilation=1, 44 | downsample=None, 45 | dimension=3): 46 | super(ECABasicBlock, self).__init__( 47 | inplanes, 48 | planes, 49 | stride=stride, 50 | dilation=dilation, 51 | downsample=downsample, 52 | dimension=dimension) 53 | self.eca = ECALayer(planes, gamma=2, b=1) 54 | 55 | def forward(self, x): 56 | residual = x 57 | 58 | out = self.conv1(x) 59 | out = self.norm1(out) 60 | out = self.relu(out) 61 | 62 | out = self.conv2(out) 63 | out = self.norm2(out) 64 | out = self.eca(out) 65 | 66 | if self.downsample is not None: 67 | residual = self.downsample(x) 68 | 69 | out += residual 70 | out = self.relu(out) 71 | 72 | return out 73 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/netvlad.py: -------------------------------------------------------------------------------- 1 | """ 2 | PointNet code taken from PointNetVLAD Pytorch implementation. 3 | """ 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.utils.data 8 | import torch.nn.functional as F 9 | import math 10 | 11 | """ 12 | NOTE: The toolbox can only pool lists of features of the same length. It was specifically optimized to efficiently 13 | do so. One way to handle multiple lists of features of variable length is to create, via a data augmentation 14 | technique, a tensor of shape: 'batch_size'x'max_samples'x'feature_size'. Where 'max_samples' would be the maximum 15 | number of feature per list. Then for each list, you would fill the tensor with 0 values. 16 | """ 17 | 18 | class NetVLADLoupe(nn.Module): 19 | def __init__(self, feature_size, cluster_size, output_dim, gating=True, add_batch_norm=True): 20 | super().__init__() 21 | self.feature_size = feature_size 22 | self.output_dim = output_dim 23 | self.gating = gating 24 | self.add_batch_norm = add_batch_norm 25 | self.cluster_size = cluster_size 26 | self.softmax = nn.Softmax(dim=-1) 27 | self.cluster_weights = nn.Parameter(torch.randn(feature_size, cluster_size) * 1 / math.sqrt(feature_size)) 28 | self.cluster_weights2 = nn.Parameter(torch.randn(1, feature_size, cluster_size) * 1 / math.sqrt(feature_size)) 29 | self.hidden1_weights = nn.Parameter( 30 | torch.randn(cluster_size * feature_size, output_dim) * 1 / math.sqrt(feature_size)) 31 | 32 | if add_batch_norm: 33 | self.cluster_biases = None 34 | self.bn1 = nn.BatchNorm1d(cluster_size) 35 | else: 36 | self.cluster_biases = nn.Parameter(torch.randn(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(output_dim, add_batch_norm=add_batch_norm) 43 | 44 | def forward(self, x): 45 | # Expects (batch_size, num_points, channels) tensor 46 | assert x.dim() == 3 47 | num_points = x.shape[1] 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, num_points, self.cluster_size) 54 | # activation = activation.transpose(1,2).contiguous() 55 | else: 56 | activation = activation + self.cluster_biases 57 | activation = self.softmax(activation) 58 | activation = activation.view((-1, num_points, self.cluster_size)) 59 | 60 | a_sum = activation.sum(-2, keepdim=True) 61 | a = a_sum * self.cluster_weights2 62 | 63 | activation = torch.transpose(activation, 2, 1) 64 | x = x.view((-1, num_points, self.feature_size)) 65 | vlad = torch.matmul(activation, x) 66 | vlad = torch.transpose(vlad, 2, 1) 67 | vlad = vlad - a 68 | 69 | vlad = F.normalize(vlad, dim=1, p=2) 70 | vlad = vlad.reshape((-1, self.cluster_size * self.feature_size)) 71 | vlad = F.normalize(vlad, dim=1, p=2) 72 | 73 | vlad = torch.matmul(vlad, self.hidden1_weights) 74 | 75 | vlad = self.bn2(vlad) 76 | 77 | if self.gating: 78 | vlad = self.context_gating(vlad) 79 | 80 | return vlad 81 | 82 | 83 | class GatingContext(nn.Module): 84 | def __init__(self, dim, add_batch_norm=True): 85 | super(GatingContext, self).__init__() 86 | self.dim = dim 87 | self.add_batch_norm = add_batch_norm 88 | self.gating_weights = nn.Parameter( 89 | torch.randn(dim, dim) * 1 / math.sqrt(dim)) 90 | self.sigmoid = nn.Sigmoid() 91 | 92 | if add_batch_norm: 93 | self.gating_biases = None 94 | self.bn1 = nn.BatchNorm1d(dim) 95 | else: 96 | self.gating_biases = nn.Parameter( 97 | torch.randn(dim) * 1 / math.sqrt(dim)) 98 | self.bn1 = None 99 | 100 | def forward(self, x): 101 | gates = torch.matmul(x, self.gating_weights) 102 | 103 | if self.add_batch_norm: 104 | gates = self.bn1(gates) 105 | else: 106 | gates = gates + self.gating_biases 107 | 108 | gates = self.sigmoid(gates) 109 | 110 | activation = x * gates 111 | 112 | return activation 113 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/pooling.py: -------------------------------------------------------------------------------- 1 | # Pooling methods code based on: https://github.com/filipradenovic/cnnimageretrieval-pytorch 2 | 3 | import torch 4 | import torch.nn as nn 5 | import MinkowskiEngine as ME 6 | 7 | from models.MinkLoc3dv2.layers.netvlad import NetVLADLoupe 8 | from torch.nn.parameter import Parameter 9 | import torch.nn.functional as F 10 | 11 | 12 | class MAC(nn.Module): 13 | def __init__(self, input_dim): 14 | super().__init__() 15 | self.input_dim = input_dim 16 | # Same output number of channels as input number of channels 17 | self.output_dim = self.input_dim 18 | self.f = ME.MinkowskiGlobalMaxPooling() 19 | 20 | def forward(self, x: ME.SparseTensor): 21 | x = self.f(x) 22 | return x.F # Return (batch_size, n_features) tensor 23 | 24 | 25 | class SPoC(nn.Module): 26 | def __init__(self, input_dim): 27 | super().__init__() 28 | self.input_dim = input_dim 29 | # Same output number of channels as input number of channels 30 | self.output_dim = self.input_dim 31 | self.f = ME.MinkowskiGlobalAvgPooling() 32 | 33 | def forward(self, x: ME.SparseTensor): 34 | x = self.f(x) 35 | return x.F # Return (batch_size, n_features) tensor 36 | 37 | 38 | class GeM(nn.Module): 39 | def __init__(self, input_dim, p=3, eps=1e-6): 40 | super(GeM, self).__init__() 41 | self.input_dim = input_dim 42 | # Same output number of channels as input number of channels 43 | self.output_dim = self.input_dim 44 | self.p = nn.Parameter(torch.ones(1) * p) 45 | self.eps = eps 46 | self.f = ME.MinkowskiGlobalAvgPooling() 47 | 48 | def forward(self, x: ME.SparseTensor): 49 | # This implicitly applies ReLU on x (clamps negative values) 50 | #temp = ME.SparseTensor(x.F.clamp(min=self.eps).pow(self.p), coordinates=x.C) 51 | 52 | temp = ME.SparseTensor(x.F.clamp(min=self.eps).pow(self.p), 53 | coordinate_manager = x.coordinate_manager, 54 | coordinate_map_key = x.coordinate_map_key) 55 | 56 | temp = self.f(temp) # Apply ME.MinkowskiGlobalAvgPooling 57 | 58 | return temp.F.pow(1./self.p) # Return (batch_size, n_features) tensor 59 | 60 | class GAP(nn.Module): 61 | def __init__(self, input_dim, p=3, eps=1e-6): 62 | super(GAP, self).__init__() 63 | self.input_dim = input_dim 64 | # Same output number of channels as input number of channels 65 | self.output_dim = self.input_dim 66 | self.f = ME.MinkowskiGlobalAvgPooling() 67 | self.eps = eps 68 | def forward(self, x: ME.SparseTensor): 69 | # This implicitly applies ReLU on x (clamps negative values) 70 | #temp = ME.SparseTensor(x.F.clamp(min=self.eps).pow(self.p), coordinates=x.C) 71 | 72 | temp = ME.SparseTensor(x.F.clamp(min=self.eps), 73 | coordinate_manager = x.coordinate_manager, 74 | coordinate_map_key = x.coordinate_map_key) 75 | 76 | temp = self.f(temp) # Apply ME.MinkowskiGlobalAvgPooling 77 | 78 | return temp.F # Return (batch_size, n_features) tensor 79 | 80 | 81 | class GeMt(nn.Module): 82 | def __init__(self, input_dim, p=1, eps=1e-6): 83 | super().__init__() 84 | self.p = Parameter(torch.ones(1)*p) 85 | self.eps = eps 86 | 87 | def gem(self, x, p=3, eps=1e-6): 88 | # return F.avg_pool2d(x.clamp(min=eps).pow(p)).pow(1./p) 89 | return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p) 90 | 91 | def flatten(self, x): 92 | assert x.shape[2] == x.shape[3] == 1, f"{x.shape[2]} != {x.shape[3]} != 1" 93 | return x[:,:,0,0] 94 | 95 | def forward(self, x): 96 | features = x.decomposed_features 97 | # features is a list of (n_points, feature_size) tensors with variable number of points 98 | features = torch.nn.utils.rnn.pad_sequence(features, batch_first=True) 99 | # features is (batch_size, n_points, feature_size) tensor padded with zeros 100 | features = features.transpose(2,1).unsqueeze(-1) 101 | x = self.gem(features, p=self.p, eps=self.eps) 102 | x = self.flatten(x) 103 | 104 | return x 105 | def __repr__(self): 106 | return self.__class__.__name__ + '(' + 'p=' + '{:.4f}'.format(self.p.data.tolist()[0]) + ', ' + 'eps=' + str(self.eps) + ')' 107 | 108 | class Avgt(nn.Module): 109 | def __init__(self, input_dim): 110 | super().__init__() 111 | self.input_dim = input_dim 112 | self.feature_size = input_dim 113 | 114 | def forward(self, x): 115 | assert x.F.shape[1] == self.feature_size 116 | features = x.decomposed_features 117 | # features is a list of (n_points, feature_size) tensors with variable number of points 118 | features = torch.nn.utils.rnn.pad_sequence(features, batch_first=True) 119 | # features is (batch_size, n_points, feature_size) tensor padded with zeros 120 | features = features.transpose(2,1) # (batch_size, feature_size, n_points) 121 | x = F.avg_pool1d(features, kernel_size= features.size(-1)).squeeze(-1) 122 | 123 | return x 124 | def __repr__(self): 125 | return self.__class__.__name__ + '(' + 'p=' + '{:.4f}'.format(self.p.data.tolist()[0]) + ', ' + 'eps=' + str(self.eps) + ')' 126 | 127 | class Avg(nn.Module): 128 | def __init__(self, input_dim, eps=1e-6): 129 | super().__init__() 130 | self.input_dim = input_dim 131 | self.feature_size = input_dim 132 | self.eps = eps 133 | self.f = ME.MinkowskiGlobalAvgPooling() 134 | 135 | 136 | def forward(self, x): 137 | 138 | assert x.F.shape[1] == self.feature_size 139 | temp = ME.SparseTensor(x.F.clamp(min=self.eps), 140 | coordinate_manager = x.coordinate_manager, 141 | coordinate_map_key = x.coordinate_map_key) 142 | 143 | x = self.f(temp) # Apply ME.MinkowskiGlobalAvgPooling 144 | 145 | 146 | # features = x.decomposed_features 147 | # # features is a list of (n_points, feature_size) tensors with variable number of points 148 | # features = torch.nn.utils.rnn.pad_sequence(features, batch_first=True) 149 | # # features is (batch_size, n_points, feature_size) tensor padded with zeros 150 | # features = features.transpose(2,1) # (batch_size, feature_size, n_points) 151 | # x = F.avg_pool1d(features, kernel_size= features.size(-1)).squeeze(-1) 152 | 153 | return x.F 154 | 155 | def __repr__(self): 156 | return self.__class__.__name__ + '(' + 'p=' + '{:.4f}'.format(self.p.data.tolist()[0]) + ', ' + 'eps=' + str(self.eps) + ')' 157 | 158 | class GeMLinear(nn.Module): 159 | def __init__(self, input_dim, out_features, p=3, eps=1e-6): 160 | super(GeMLinear, self).__init__() 161 | self.input_dim = input_dim 162 | # Same output number of channels as input number of channels 163 | self.output_dim = out_features 164 | self.p = p 165 | self.eps = eps 166 | self.p = nn.Parameter(torch.ones(1) * p) 167 | self.f = ME.MinkowskiGlobalAvgPooling() 168 | 169 | def forward(self, x: ME.SparseTensor): 170 | # This implicitly applies ReLU on x (clamps negative values) 171 | #temp = ME.SparseTensor(x.F.clamp(min=self.eps).pow(self.p), coordinates=x.C) 172 | temp = ME.SparseTensor(x.F.clamp(min=self.eps).pow(self.p), 173 | coordinate_manager = x.coordinate_manager, 174 | coordinate_map_key = x.coordinate_map_key) 175 | temp = self.f(temp) # Apply ME.MinkowskiGlobalAvgPooling 176 | 177 | temp = ME.SparseTensor(temp.F.pow(1./self.p), 178 | coordinate_manager = x.coordinate_manager, 179 | coordinate_map_key = x.coordinate_map_key) 180 | 181 | return temp # Return (batch_size, n_features) tensor 182 | 183 | 184 | class NetVLADWrapper(nn.Module): 185 | def __init__(self, feature_size, output_dim, gating=True): 186 | super().__init__() 187 | self.feature_size = feature_size 188 | self.output_dim = output_dim 189 | self.net_vlad = NetVLADLoupe(feature_size=feature_size, cluster_size=64, output_dim=output_dim, gating=gating, 190 | add_batch_norm=True) 191 | 192 | def forward(self, x: ME.SparseTensor): 193 | # x is (batch_size, C, H, W) 194 | assert x.F.shape[1] == self.feature_size 195 | 196 | features = x.decomposed_features 197 | # features is a list of (n_points, feature_size) tensors with variable number of points 198 | batch_size = len(features) 199 | features = torch.nn.utils.rnn.pad_sequence(features, batch_first=True) 200 | # features is (batch_size, n_points, feature_size) tensor padded with zeros 201 | 202 | x = self.net_vlad(features) 203 | assert x.shape[0] == batch_size 204 | assert x.shape[1] == self.output_dim 205 | return x # Return (batch_size, output_dim) tensor 206 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/layers/pooling_wrapper.py: -------------------------------------------------------------------------------- 1 | from models.MinkLoc3dv2.layers.pooling import MAC, SPoC, GeM, NetVLADWrapper, \ 2 | GeMt, Avg, Avgt, GeMLinear, GAP 3 | import torch.nn as nn 4 | import MinkowskiEngine as ME 5 | 6 | 7 | class PoolingWrapper(nn.Module): 8 | def __init__(self, pool_method, in_dim, output_dim): 9 | super().__init__() 10 | 11 | self.pool_method = pool_method 12 | self.in_dim = in_dim 13 | self.output_dim = output_dim 14 | 15 | if pool_method == 'MAC': 16 | # Global max pooling 17 | assert in_dim == output_dim 18 | self.pooling = MAC(input_dim=in_dim) 19 | elif pool_method == 'SPoC': 20 | # Global average pooling 21 | assert in_dim == output_dim 22 | self.pooling = SPoC(input_dim=in_dim) 23 | elif pool_method == 'GeM': 24 | # Generalized mean pooling 25 | assert in_dim == output_dim 26 | self.pooling = GeM(input_dim=in_dim) 27 | elif pool_method == 'GAP': 28 | # Generalized mean pooling 29 | assert in_dim == output_dim 30 | self.pooling = GAP(input_dim=in_dim) 31 | elif pool_method == 'GeMt': 32 | # Generalized mean pooling 33 | assert in_dim == output_dim 34 | self.pooling = GeMt(input_dim=in_dim) 35 | elif pool_method == 'Avgt': 36 | # Generalized mean pooling 37 | self.pooling = Avgt(input_dim=in_dim) 38 | elif pool_method == 'Avg': 39 | # Generalized mean pooling 40 | self.pooling = Avg(input_dim=in_dim) 41 | elif pool_method == 'GeMLinear': 42 | # Generalized mean pooling 43 | self.pooling = GeMLinear(input_dim=in_dim, out_features = output_dim) 44 | elif self.pool_method == 'netvlad': 45 | # NetVLAD 46 | self.pooling = NetVLADWrapper(feature_size=in_dim, output_dim=output_dim, gating=False) 47 | elif self.pool_method == 'netvladgc': 48 | # NetVLAD with Gating Context 49 | self.pooling = NetVLADWrapper(feature_size=in_dim, output_dim=output_dim, gating=True) 50 | else: 51 | raise NotImplementedError('Unknown pooling method: {}'.format(pool_method)) 52 | 53 | def forward(self, x: ME.SparseTensor): 54 | return self.pooling(x) 55 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/losses/__pycache__/loss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/losses/__pycache__/loss.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/losses/__pycache__/loss_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/losses/__pycache__/loss_utils.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/losses/__pycache__/truncated_smoothap.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/MinkLoc3dv2/losses/__pycache__/truncated_smoothap.cpython-38.pyc -------------------------------------------------------------------------------- /models/MinkLoc3dv2/losses/loss.py: -------------------------------------------------------------------------------- 1 | # Warsaw University of Technology 2 | 3 | from pytorch_metric_learning import losses, reducers 4 | from pytorch_metric_learning.distances import LpDistance 5 | from models.MinkLoc3dv2.mink_params import TrainingParams 6 | from models.MinkLoc3dv2.losses.loss_utils import * 7 | from models.MinkLoc3dv2.losses.truncated_smoothap import TruncatedSmoothAP 8 | 9 | 10 | def make_losses(params: TrainingParams): 11 | if params.loss == 'batchhardtripletmarginloss': 12 | # BatchHard mining with triplet margin loss 13 | # Expects input: embeddings, positives_mask, negatives_mask 14 | loss_fn = BatchHardTripletLossWithMasks(params.margin) 15 | elif params.loss == 'batchhardcontrastiveloss': 16 | loss_fn = BatchHardContrastiveLossWithMasks(params.pos_margin, params.neg_margin) 17 | elif params.loss == 'truncatedsmoothap': 18 | loss_fn = TruncatedSmoothAP(tau1=params.tau1, similarity=params.similarity, 19 | positives_per_query=params.positives_per_query) 20 | else: 21 | print('Unknown loss: {}'.format(params.loss)) 22 | raise NotImplementedError 23 | 24 | return loss_fn 25 | 26 | 27 | class HardTripletMinerWithMasks: 28 | # Hard triplet miner 29 | def __init__(self, distance): 30 | self.distance = distance 31 | # Stats 32 | self.max_pos_pair_dist = None 33 | self.max_neg_pair_dist = None 34 | self.mean_pos_pair_dist = None 35 | self.mean_neg_pair_dist = None 36 | self.min_pos_pair_dist = None 37 | self.min_neg_pair_dist = None 38 | 39 | def __call__(self, embeddings, positives_mask, negatives_mask): 40 | assert embeddings.dim() == 2 41 | d_embeddings = embeddings.detach() 42 | with torch.no_grad(): 43 | hard_triplets = self.mine(d_embeddings, positives_mask, negatives_mask) 44 | return hard_triplets 45 | 46 | def mine(self, embeddings, positives_mask, negatives_mask): 47 | # Based on pytorch-metric-learning implementation 48 | dist_mat = self.distance(embeddings) 49 | (hardest_positive_dist, hardest_positive_indices), a1p_keep = get_max_per_row(dist_mat, positives_mask) 50 | (hardest_negative_dist, hardest_negative_indices), a2n_keep = get_min_per_row(dist_mat, negatives_mask) 51 | a_keep_idx = torch.where(a1p_keep & a2n_keep) 52 | a = torch.arange(dist_mat.size(0)).to(hardest_positive_indices.device)[a_keep_idx] 53 | p = hardest_positive_indices[a_keep_idx] 54 | n = hardest_negative_indices[a_keep_idx] 55 | self.max_pos_pair_dist = torch.max(hardest_positive_dist[a_keep_idx]).item() 56 | self.max_neg_pair_dist = torch.max(hardest_negative_dist[a_keep_idx]).item() 57 | self.mean_pos_pair_dist = torch.mean(hardest_positive_dist[a_keep_idx]).item() 58 | self.mean_neg_pair_dist = torch.mean(hardest_negative_dist[a_keep_idx]).item() 59 | self.min_pos_pair_dist = torch.min(hardest_positive_dist[a_keep_idx]).item() 60 | self.min_neg_pair_dist = torch.min(hardest_negative_dist[a_keep_idx]).item() 61 | return a, p, n 62 | 63 | 64 | def get_max_per_row(mat, mask): 65 | non_zero_rows = torch.any(mask, dim=1) 66 | mat_masked = mat.clone() 67 | mat_masked[~mask] = 0 68 | return torch.max(mat_masked, dim=1), non_zero_rows 69 | 70 | 71 | def get_min_per_row(mat, mask): 72 | non_inf_rows = torch.any(mask, dim=1) 73 | mat_masked = mat.clone() 74 | mat_masked[~mask] = float('inf') 75 | return torch.min(mat_masked, dim=1), non_inf_rows 76 | 77 | 78 | class BatchHardTripletLossWithMasks: 79 | def __init__(self, margin:float): 80 | self.margin = margin 81 | self.distance = LpDistance(normalize_embeddings=False, collect_stats=True) 82 | # We use triplet loss with Euclidean distance 83 | self.miner_fn = HardTripletMinerWithMasks(distance=self.distance) 84 | reducer_fn = reducers.AvgNonZeroReducer(collect_stats=True) 85 | self.loss_fn = losses.TripletMarginLoss(margin=self.margin, swap=True, distance=self.distance, 86 | reducer=reducer_fn, collect_stats=True) 87 | 88 | def __call__(self, embeddings, positives_mask, negatives_mask): 89 | hard_triplets = self.miner_fn(embeddings, positives_mask, negatives_mask) 90 | dummy_labels = torch.arange(embeddings.shape[0]).to(embeddings.device) 91 | loss = self.loss_fn(embeddings, dummy_labels, hard_triplets) 92 | 93 | stats = {'loss': loss.item(), 'avg_embedding_norm': self.loss_fn.distance.final_avg_query_norm, 94 | 'num_non_zero_triplets': self.loss_fn.reducer.triplets_past_filter, 95 | 'num_triplets': len(hard_triplets[0]), 96 | 'mean_pos_pair_dist': self.miner_fn.mean_pos_pair_dist, 97 | 'mean_neg_pair_dist': self.miner_fn.mean_neg_pair_dist, 98 | 'max_pos_pair_dist': self.miner_fn.max_pos_pair_dist, 99 | 'max_neg_pair_dist': self.miner_fn.max_neg_pair_dist, 100 | 'min_pos_pair_dist': self.miner_fn.min_pos_pair_dist, 101 | 'min_neg_pair_dist': self.miner_fn.min_neg_pair_dist 102 | } 103 | return loss, stats 104 | 105 | 106 | class BatchHardContrastiveLossWithMasks: 107 | def __init__(self, pos_margin: float, neg_margin: float): 108 | self.pos_margin = pos_margin 109 | self.neg_margin = neg_margin 110 | self.distance = LpDistance(normalize_embeddings=False, collect_stats=True) 111 | self.miner_fn = HardTripletMinerWithMasks(distance=self.distance) 112 | # We use contrastive loss with squared Euclidean distance 113 | reducer_fn = reducers.AvgNonZeroReducer(collect_stats=True) 114 | self.loss_fn = losses.ContrastiveLoss(pos_margin=self.pos_margin, neg_margin=self.neg_margin, 115 | distance=self.distance, reducer=reducer_fn, collect_stats=True) 116 | 117 | def __call__(self, embeddings, positives_mask, negatives_mask): 118 | hard_triplets = self.miner_fn(embeddings, positives_mask, negatives_mask) 119 | dummy_labels = torch.arange(embeddings.shape[0]).to(embeddings.device) 120 | loss = self.loss_fn(embeddings, dummy_labels, hard_triplets) 121 | stats = {'loss': loss.item(), 'avg_embedding_norm': self.loss_fn.distance.final_avg_query_norm, 122 | 'pos_pairs_above_threshold': self.loss_fn.reducer.reducers['pos_loss'].pos_pairs_above_threshold, 123 | 'neg_pairs_above_threshold': self.loss_fn.reducer.reducers['neg_loss'].neg_pairs_above_threshold, 124 | 'pos_loss': self.loss_fn.reducer.reducers['pos_loss'].pos_loss.item(), 125 | 'neg_loss': self.loss_fn.reducer.reducers['neg_loss'].neg_loss.item(), 126 | 'num_pairs': 2*len(hard_triplets[0]), 127 | 'mean_pos_pair_dist': self.miner_fn.mean_pos_pair_dist, 128 | 'mean_neg_pair_dist': self.miner_fn.mean_neg_pair_dist, 129 | 'max_pos_pair_dist': self.miner_fn.max_pos_pair_dist, 130 | 'max_neg_pair_dist': self.miner_fn.max_neg_pair_dist, 131 | 'min_pos_pair_dist': self.miner_fn.min_pos_pair_dist, 132 | 'min_neg_pair_dist': self.miner_fn.min_neg_pair_dist 133 | } 134 | 135 | return loss, stats 136 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/losses/loss_utils.py: -------------------------------------------------------------------------------- 1 | # Functions and classes used by different loss functions 2 | import numpy as np 3 | import torch 4 | from torch import Tensor 5 | 6 | EPS = 1e-5 7 | 8 | 9 | def metrics_mean(l): 10 | # Compute the mean and return as Python number 11 | metrics = {} 12 | for e in l: 13 | for metric_name in e: 14 | if metric_name not in metrics: 15 | metrics[metric_name] = [] 16 | metrics[metric_name].append(e[metric_name]) 17 | 18 | for metric_name in metrics: 19 | metrics[metric_name] = np.mean(np.array(metrics[metric_name])) 20 | 21 | return metrics 22 | 23 | 24 | def squared_euclidean_distance(x: Tensor, y: Tensor) -> Tensor: 25 | ''' 26 | Compute squared Euclidean distance 27 | Input: x is Nxd matrix 28 | y is Mxd matirx 29 | Output: dist is a NxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:] 30 | i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2 31 | Source: https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/3 32 | ''' 33 | x_norm = (x ** 2).sum(1).view(-1, 1) 34 | y_t = torch.transpose(y, 0, 1) 35 | y_norm = (y ** 2).sum(1).view(1, -1) 36 | dist = x_norm + y_norm - 2.0 * torch.mm(x, y_t) 37 | return torch.clamp(dist, 0.0, np.inf) 38 | 39 | 40 | def sigmoid(tensor: Tensor, temp: float) -> Tensor: 41 | """ temperature controlled sigmoid 42 | takes as input a torch tensor (tensor) and passes it through a sigmoid, controlled by temperature: temp 43 | """ 44 | exponent = -tensor / temp 45 | # clamp the input tensor for stability 46 | exponent = torch.clamp(exponent, min=-50, max=50) 47 | y = 1.0 / (1.0 + torch.exp(exponent)) 48 | return y 49 | 50 | 51 | def compute_aff(x: Tensor, similarity: str = 'cosine') -> Tensor: 52 | """computes the affinity matrix between an input vector and itself""" 53 | if similarity == 'cosine': 54 | x = torch.mm(x, x.t()) 55 | elif similarity == 'euclidean': 56 | x = x.unsqueeze(0) 57 | x = torch.cdist(x, x, p=2) 58 | x = x.squeeze(0) 59 | # The greater the distance the smaller affinity 60 | x = -x 61 | else: 62 | raise NotImplementedError(f"Incorrect similarity measure: {similarity}") 63 | return x -------------------------------------------------------------------------------- /models/MinkLoc3dv2/losses/truncated_smoothap.py: -------------------------------------------------------------------------------- 1 | # Implemented as per "Recall@k Surrogate Loss with Large Batches and Similarity Mixup" paper 2 | # but only the fixed number of the closest positives is considered 3 | 4 | import numpy as np 5 | import torch 6 | 7 | from models.MinkLoc3dv2.losses.loss_utils import sigmoid, compute_aff 8 | 9 | 10 | class TruncatedSmoothAP: 11 | def __init__(self, tau1: float = 0.01, similarity: str = 'cosine', positives_per_query: int = 4): 12 | # We reversed the notation compared to the paper (tau1 is sigmoid on similarity differences) 13 | # tau1: sigmoid temperature applied on similarity differences 14 | # positives_per_query: number of positives per query to consider 15 | # negatives_only: if True in denominator we consider positives and negatives; if False we consider all elements 16 | # (with except to the anchor itself) 17 | 18 | self.tau1 = tau1 19 | self.similarity = similarity 20 | self.positives_per_query = positives_per_query 21 | 22 | def __call__(self, embeddings, positives_mask, negatives_mask): 23 | device = embeddings.device 24 | 25 | positives_mask = positives_mask.to(device) 26 | negatives_mask = negatives_mask.to(device) 27 | 28 | # Ranking of the retrieval set 29 | # For each element we ignore elements that are neither positives nor negatives 30 | 31 | # Compute cosine similarity scores 32 | # 1st dimension corresponds to q, 2nd dimension to z 33 | s_qz = compute_aff(embeddings, similarity=self.similarity) 34 | 35 | # Find the positives_per_query closest positives for each query 36 | s_positives = s_qz.detach().clone() 37 | s_positives.masked_fill_(torch.logical_not(positives_mask), np.NINF) 38 | #closest_positives_ndx = torch.argmax(s_positives, dim=1).view(-1, 1) # Indices of closests positives for each query 39 | closest_positives_ndx = torch.topk(s_positives, k=self.positives_per_query, dim=1, largest=True, sorted=True)[1] 40 | # closest_positives_ndx is (batch_size, positives_per_query) with positives_per_query closest positives 41 | # per each batch element 42 | 43 | n_positives = positives_mask.sum(dim=1) # Number of positives for each anchor 44 | 45 | # Compute the rank of each example x with respect to query element q as per Eq. (2) 46 | s_diff = s_qz.unsqueeze(1) - s_qz.gather(1, closest_positives_ndx).unsqueeze(2) 47 | s_sigmoid = sigmoid(s_diff, temp=self.tau1) 48 | 49 | # Compute the nominator in Eq. 2 and 5 - for q compute the ranking of each of its positives with respect to other positives of q 50 | # Filter out z not in Positives 51 | pos_mask = positives_mask.unsqueeze(1) 52 | pos_s_sigmoid = s_sigmoid * pos_mask 53 | 54 | # Filter out z on the same position as the positive (they have value = 0.5, as the similarity difference is zero) 55 | mask = torch.ones_like(pos_s_sigmoid).scatter(2, closest_positives_ndx.unsqueeze(2), 0.) 56 | pos_s_sigmoid = pos_s_sigmoid * mask 57 | 58 | # Compute the rank for each query and its positives_per_query closest positive examples with respect to other positives 59 | r_p = torch.sum(pos_s_sigmoid, dim=2) + 1. 60 | # r_p is (batch_size, positives_per_query) matrix 61 | 62 | # Consider only positives and negatives in the denominator 63 | # Compute the denominator in Eq. 5 - add sum of Indicator function for negatives (or non-positives) 64 | neg_mask = negatives_mask.unsqueeze(1) 65 | neg_s_sigmoid = s_sigmoid * neg_mask 66 | r_omega = r_p + torch.sum(neg_s_sigmoid, dim=2) 67 | 68 | # Compute R(i, S_p) / R(i, S_omega) ration in Eq. 2 69 | r = r_p / r_omega 70 | 71 | # Compute metrics mean ranking of the positive example, recall@1 72 | stats = {} 73 | # Mean number of positives per query 74 | stats['positives_per_query'] = n_positives.float().mean(dim=0).item() 75 | # Mean ranking of selected positive examples (closests positives) 76 | temp = s_diff.detach() > 0 77 | temp = torch.logical_and(temp[:, 0], negatives_mask) # Take the best positive 78 | hard_ranking = temp.sum(dim=1) 79 | stats['best_positive_ranking'] = hard_ranking.float().mean(dim=0).item() 80 | # Recall at 1 81 | stats['recall'] = {1: (hard_ranking <= 1).float().mean(dim=0).item()} 82 | 83 | # r is (N, positives_per_query) tensor 84 | # Zero entries not corresponding to real positives - this happens when the number of true positives is lower than positives_per_query 85 | valid_positives_mask = torch.gather(positives_mask, 1, closest_positives_ndx) # () tensor 86 | masked_r = r * valid_positives_mask 87 | n_valid_positives = valid_positives_mask.sum(dim=1) 88 | 89 | # Filter out rows (queries) without any positive to avoid division by zero 90 | valid_q_mask = n_valid_positives > 0 91 | masked_r = masked_r[valid_q_mask] 92 | 93 | ap = (masked_r.sum(dim=1) / n_valid_positives[valid_q_mask]).mean() 94 | loss = 1. - ap 95 | 96 | stats['loss'] = loss.item() 97 | stats['ap'] = ap.item() 98 | stats['avg_embedding_norm'] = embeddings.norm(dim=1).mean().item() 99 | return loss, stats 100 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/minkfpn.py: -------------------------------------------------------------------------------- 1 | # Warsaw University of Technology 2 | 3 | import torch.nn as nn 4 | import MinkowskiEngine as ME 5 | from MinkowskiEngine.modules.resnet_block import BasicBlock 6 | from models.MinkLoc3dv2.resnet import ResNetBase 7 | 8 | 9 | class MinkFPN(ResNetBase): 10 | # Feature Pyramid Network (FPN) architecture implementation using Minkowski ResNet building blocks 11 | def __init__(self, in_channels, out_channels, num_top_down=1, conv0_kernel_size=5, block=BasicBlock, 12 | layers=(1, 1, 1), planes=(32, 64, 64)): 13 | assert len(layers) == len(planes) 14 | assert 1 <= len(layers) 15 | assert 0 <= num_top_down <= len(layers) 16 | self.num_bottom_up = len(layers) 17 | self.num_top_down = num_top_down 18 | self.conv0_kernel_size = conv0_kernel_size 19 | self.block = block 20 | self.layers = layers 21 | self.planes = planes 22 | self.lateral_dim = out_channels 23 | self.init_dim = planes[0] 24 | ResNetBase.__init__(self, in_channels, out_channels, D=3) 25 | 26 | def network_initialization(self, in_channels, out_channels, D): 27 | assert len(self.layers) == len(self.planes) 28 | assert len(self.planes) == self.num_bottom_up 29 | 30 | self.convs = nn.ModuleList() # Bottom-up convolutional blocks with stride=2 31 | self.bn = nn.ModuleList() # Bottom-up BatchNorms 32 | self.blocks = nn.ModuleList() # Bottom-up blocks 33 | self.tconvs = nn.ModuleList() # Top-down tranposed convolutions 34 | self.conv1x1 = nn.ModuleList() # 1x1 convolutions in lateral connections 35 | 36 | # The first convolution is special case, with kernel size = 5 37 | self.inplanes = self.planes[0] 38 | self.conv0 = ME.MinkowskiConvolution(in_channels, self.inplanes, kernel_size=self.conv0_kernel_size, 39 | dimension=D) 40 | self.bn0 = ME.MinkowskiBatchNorm(self.inplanes) 41 | 42 | for plane, layer in zip(self.planes, self.layers): 43 | self.convs.append(ME.MinkowskiConvolution(self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)) 44 | self.bn.append(ME.MinkowskiBatchNorm(self.inplanes)) 45 | self.blocks.append(self._make_layer(self.block, plane, layer)) 46 | 47 | # Lateral connections 48 | for i in range(self.num_top_down): 49 | self.conv1x1.append(ME.MinkowskiConvolution(self.planes[-1 - i], self.lateral_dim, kernel_size=1, 50 | stride=1, dimension=D)) 51 | self.tconvs.append(ME.MinkowskiConvolutionTranspose(self.lateral_dim, self.lateral_dim, kernel_size=2, 52 | stride=2, dimension=D)) 53 | # There's one more lateral connection than top-down TConv blocks 54 | if self.num_top_down < self.num_bottom_up: 55 | # Lateral connection from Conv block 1 or above 56 | self.conv1x1.append(ME.MinkowskiConvolution(self.planes[-1 - self.num_top_down], self.lateral_dim, kernel_size=1, 57 | stride=1, dimension=D)) 58 | else: 59 | # Lateral connection from Con0 block 60 | self.conv1x1.append(ME.MinkowskiConvolution(self.planes[0], self.lateral_dim, kernel_size=1, 61 | stride=1, dimension=D)) 62 | 63 | self.relu = ME.MinkowskiReLU(inplace=True) 64 | 65 | def forward(self, x): 66 | # *** BOTTOM-UP PASS *** 67 | # First bottom-up convolution is special (with bigger kernel) 68 | feature_maps = [] 69 | x = self.conv0(x) 70 | x = self.bn0(x) 71 | x = self.relu(x) 72 | if self.num_top_down == self.num_bottom_up: 73 | feature_maps.append(x) 74 | 75 | # BOTTOM-UP PASS 76 | for ndx, (conv, bn, block) in enumerate(zip(self.convs, self.bn, self.blocks)): 77 | x = conv(x) # Downsample (conv stride=2 with 2x2x2 kernel) 78 | x = bn(x) 79 | x = self.relu(x) 80 | x = block(x) 81 | if self.num_bottom_up - 1 - self.num_top_down <= ndx < len(self.convs) - 1: 82 | feature_maps.append(x) 83 | 84 | assert len(feature_maps) == self.num_top_down 85 | 86 | x = self.conv1x1[0](x) 87 | 88 | # TOP-DOWN PASS 89 | for ndx, tconv in enumerate(self.tconvs): 90 | x = tconv(x) # Upsample using transposed convolution 91 | x = x + self.conv1x1[ndx+1](feature_maps[-ndx - 1]) 92 | 93 | return x 94 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/minkloc.py: -------------------------------------------------------------------------------- 1 | # Author: Jacek Komorowski 2 | # Warsaw University of Technology 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | import MinkowskiEngine as ME 8 | 9 | from models.MinkLoc3dv2.layers.pooling_wrapper import PoolingWrapper 10 | 11 | 12 | class MinkLoc(torch.nn.Module): 13 | def __init__(self, backbone: nn.Module, pooling: PoolingWrapper, normalize_embeddings: bool = False): 14 | super().__init__() 15 | self.backbone = backbone 16 | self.pooling = pooling 17 | self.normalize_embeddings = normalize_embeddings 18 | self.stats = {} 19 | self.normalize_embeddings = normalize_embeddings 20 | 21 | def update_aggregators(self): 22 | print('No update!') 23 | 24 | def forward(self, batch, is_training =True): 25 | x = ME.SparseTensor(batch['features'], coordinates=batch['coords']) 26 | 27 | x = self.backbone(x) 28 | # x is (num_points, n_features) tensor 29 | assert x.shape[1] == self.pooling.in_dim, f'Backbone output tensor has: {x.shape[1]} channels. ' \ 30 | f'Expected: {self.pooling.in_dim}' 31 | 32 | x = self.pooling(x) 33 | if self.normalize_embeddings: 34 | x = F.normalize(x, p=2, dim=1) 35 | 36 | if hasattr(self.pooling, 'stats'): 37 | self.stats.update(self.pooling.stats) 38 | #x = x.flatten(1) 39 | assert x.dim() == 2, f'Expected 2-dimensional tensor (batch_size,output_dim). Got {x.dim()} dimensions.' 40 | # assert x.shape[1] == self.pooling.output_dim, f'Output tensor has: {x.shape[1]} channels. ' \ 41 | # f'Expected: {self.pooling.output_dim}' 42 | 43 | return x 44 | 45 | def print_info(self): 46 | print('Model class: MinkLoc') 47 | n_params = sum([param.nelement() for param in self.parameters()]) 48 | print(f'Total parameters: {n_params}') 49 | n_params = sum([param.nelement() for param in self.backbone.parameters()]) 50 | print(f'Backbone: {type(self.backbone).__name__} #parameters: {n_params}') 51 | n_params = sum([param.nelement() for param in self.pooling.parameters()]) 52 | print(f'Pooling method: {self.pooling.pool_method} #parameters: {n_params}') 53 | print('# channels from the backbone: {}'.format(self.pooling.in_dim)) 54 | print('# output channels : {}'.format(self.pooling.output_dim)) 55 | print(f'Embedding normalization: {self.normalize_embeddings}') 56 | 57 | 58 | class MinkLocLAWS(torch.nn.Module): 59 | def __init__(self, backbone: nn.Module, 60 | pool_method: str = 'GeM', 61 | in_dim: int = 256, 62 | output_dim: int = 256, 63 | normalize_embeddings: bool = False): 64 | super().__init__() 65 | self.backbone = backbone 66 | self.pool_method = pool_method 67 | self.in_dim = in_dim 68 | self.output_dim = output_dim 69 | self.normalize_embeddings = normalize_embeddings 70 | self.stats = {} 71 | self.aggregators = nn.ModuleList() 72 | self.curr_group = 0 73 | print('self.normalize_embeddings', self.normalize_embeddings) 74 | 75 | def update_aggregators(self): 76 | pooling = PoolingWrapper(pool_method=self.pool_method, in_dim=self.in_dim, 77 | output_dim=self.output_dim) 78 | self.aggregators.append(pooling) 79 | 80 | def feature_extraction_training_module(self, x): 81 | x = self.aggregators[self.curr_group](x) 82 | if self.normalize_embeddings: 83 | x = F.normalize(x, p=2, dim=1) 84 | return x 85 | 86 | def feature_extraction_inference_module(self, x): 87 | global_features = [] 88 | for aggregator in self.aggregators: 89 | g_des = aggregator(x) 90 | if self.normalize_embeddings: 91 | g_des = F.normalize(g_des, p=2, dim=1) 92 | global_features.append(g_des) 93 | global_features = torch.cat(global_features, 1) 94 | return global_features 95 | 96 | def forward(self, batch, is_training =True): 97 | x = ME.SparseTensor(batch['features'], coordinates=batch['coords']) 98 | 99 | x = self.backbone(x) 100 | # x is (num_points, n_features) tensor 101 | assert x.shape[1] == self.aggregators[self.curr_group].in_dim, f'Backbone output tensor has: {x.shape[1]} channels. ' \ 102 | f'Expected: {self.pooling.in_dim}' 103 | 104 | if is_training: 105 | x = self.feature_extraction_training_module(x) 106 | else: 107 | x = self.feature_extraction_inference_module(x) 108 | 109 | if hasattr(self.aggregators[self.curr_group], 'stats'): 110 | self.stats.update(self.aggregators[self.curr_group].stats) 111 | #x = x.flatten(1) 112 | assert x.dim() == 2, f'Expected 2-dimensional tensor (batch_size,output_dim). Got {x.dim()} dimensions.' 113 | # assert x.shape[1] == self.aggregators[self.curr_group].output_dim, f'Output tensor has: {x.shape[1]} channels. ' \ 114 | # f'Expected: {self.pooling.output_dim}' 115 | 116 | return x 117 | 118 | def print_info(self): 119 | print('Model class: MinkLoc') 120 | n_params = sum([param.nelement() for param in self.parameters()]) 121 | print(f'Total parameters: {n_params}') 122 | n_params = sum([param.nelement() for param in self.backbone.parameters()]) 123 | print(f'Backbone: {type(self.backbone).__name__} #parameters: {n_params}') 124 | n_params = sum([param.nelement() for param in self.pooling.parameters()]) 125 | print(f'Pooling method: {self.pooling.pool_method} #parameters: {n_params}') 126 | print('# channels from the backbone: {}'.format(self.pooling.in_dim)) 127 | print('# output channels : {}'.format(self.pooling.output_dim)) 128 | print(f'Embedding normalization: {self.normalize_embeddings}') 129 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/minkloc3dv2.txt: -------------------------------------------------------------------------------- 1 | [MODEL] 2 | model=mink 3 | planes=64,128,64,32 4 | layers=1,1,1,1 5 | num_top_down=2 6 | conv0_kernel_size=5 7 | feature_size=256 8 | block=ECABasicBlock 9 | pooling=netvlad 10 | 11 | coordinates=cartesian 12 | quantization_step=0.01 13 | 14 | normalize_embeddings=False 15 | 16 | -------------------------------------------------------------------------------- /models/MinkLoc3dv2/resnet.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Chris Choy (chrischoy@ai.stanford.edu). 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 4 | # this software and associated documentation files (the "Software"), to deal in 5 | # the Software without restriction, including without limitation the rights to 6 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 7 | # of the Software, and to permit persons to whom the Software is furnished to do 8 | # so, subject to the following conditions: 9 | # 10 | # The above copyright notice and this permission notice shall be included in all 11 | # copies or substantial portions of the Software. 12 | # 13 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 19 | # SOFTWARE. 20 | # 21 | # Please cite "4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural 22 | # Networks", CVPR'19 (https://arxiv.org/abs/1904.08755) if you use any part 23 | # of the code. 24 | 25 | import torch.nn as nn 26 | 27 | import MinkowskiEngine as ME 28 | from MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck 29 | 30 | 31 | class ResNetBase(nn.Module): 32 | block = None 33 | layers = () 34 | init_dim = 64 35 | planes = (64, 128, 256, 512) 36 | 37 | def __init__(self, in_channels, out_channels, D=3): 38 | nn.Module.__init__(self) 39 | self.D = D 40 | assert self.block is not None 41 | 42 | self.network_initialization(in_channels, out_channels, D) 43 | self.weight_initialization() 44 | 45 | def network_initialization(self, in_channels, out_channels, D): 46 | self.inplanes = self.init_dim 47 | self.conv1 = ME.MinkowskiConvolution(in_channels, self.inplanes, kernel_size=5, stride=2, dimension=D) 48 | 49 | self.bn1 = ME.MinkowskiBatchNorm(self.inplanes) 50 | self.relu = ME.MinkowskiReLU(inplace=True) 51 | 52 | self.pool = ME.MinkowskiAvgPooling(kernel_size=2, stride=2, dimension=D) 53 | 54 | self.layer1 = self._make_layer(self.block, self.planes[0], self.layers[0], stride=2) 55 | self.layer2 = self._make_layer(self.block, self.planes[1], self.layers[1], stride=2) 56 | self.layer3 = self._make_layer(self.block, self.planes[2], self.layers[2], stride=2) 57 | self.layer4 = self._make_layer(self.block, self.planes[3], self.layers[3], stride=2) 58 | 59 | self.conv5 = ME.MinkowskiConvolution(self.inplanes, self.inplanes, kernel_size=3, stride=3, dimension=D) 60 | self.bn5 = ME.MinkowskiBatchNorm(self.inplanes) 61 | self.glob_avg = ME.MinkowskiGlobalMaxPooling() 62 | self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True) 63 | 64 | def weight_initialization(self): 65 | for m in self.modules(): 66 | if isinstance(m, ME.MinkowskiConvolution): 67 | ME.utils.kaiming_normal_(m.kernel, mode='fan_out', nonlinearity='relu') 68 | 69 | if isinstance(m, ME.MinkowskiBatchNorm): 70 | nn.init.constant_(m.bn.weight, 1) 71 | nn.init.constant_(m.bn.bias, 0) 72 | 73 | def _make_layer(self, block, planes, blocks, stride=1, dilation=1, bn_momentum=0.1): 74 | downsample = None 75 | if stride != 1 or self.inplanes != planes * block.expansion: 76 | downsample = nn.Sequential(ME.MinkowskiConvolution(self.inplanes, planes * block.expansion, 77 | kernel_size=1, stride=stride, dimension=self.D), 78 | ME.MinkowskiBatchNorm(planes * block.expansion)) 79 | layers = [] 80 | layers.append(block(self.inplanes, planes, stride=stride, dilation=dilation, downsample=downsample, 81 | dimension=self.D)) 82 | self.inplanes = planes * block.expansion 83 | for i in range(1, blocks): 84 | layers.append(block(self.inplanes, planes, stride=1, dilation=dilation, dimension=self.D)) 85 | 86 | return nn.Sequential(*layers) 87 | 88 | def forward(self, x): 89 | x = self.conv1(x) 90 | x = self.bn1(x) 91 | x = self.relu(x) 92 | x = self.pool(x) 93 | 94 | x = self.layer1(x) 95 | x = self.layer2(x) 96 | x = self.layer3(x) 97 | x = self.layer4(x) 98 | 99 | x = self.conv5(x) 100 | x = self.bn5(x) 101 | x = self.relu(x) 102 | 103 | x = self.glob_avg(x) 104 | return self.final(x) 105 | 106 | 107 | class ResNet14(ResNetBase): 108 | BLOCK = BasicBlock 109 | LAYERS = (1, 1, 1, 1) 110 | 111 | 112 | class ResNet18(ResNetBase): 113 | BLOCK = BasicBlock 114 | LAYERS = (2, 2, 2, 2) 115 | 116 | 117 | class ResNet34(ResNetBase): 118 | BLOCK = BasicBlock 119 | LAYERS = (3, 4, 6, 3) 120 | 121 | 122 | class ResNet50(ResNetBase): 123 | BLOCK = Bottleneck 124 | LAYERS = (3, 4, 6, 3) 125 | 126 | 127 | class ResNet101(ResNetBase): 128 | BLOCK = Bottleneck 129 | LAYERS = (3, 4, 23, 3) 130 | -------------------------------------------------------------------------------- /models/__pycache__/PointNetVlad.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/__pycache__/PointNetVlad.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/loupe.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/__pycache__/loupe.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/model_factory.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/__pycache__/model_factory.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/pptnet_v2.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/__pycache__/pptnet_v2.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/pt_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/models/__pycache__/pt_util.cpython-38.pyc -------------------------------------------------------------------------------- /models/loupe.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | class NetVLADBase(nn.Module): 7 | def __init__(self, feature_size, max_samples, cluster_size, output_dim, 8 | gating=True, add_batch_norm=True): 9 | super(NetVLADBase, self).__init__() 10 | self.feature_size = feature_size 11 | self.max_samples = max_samples 12 | self.output_dim = output_dim 13 | self.gating = gating 14 | self.add_batch_norm = add_batch_norm 15 | self.cluster_size = cluster_size 16 | self.softmax = nn.Softmax(dim=-1) 17 | 18 | self.cluster_weights = nn.Parameter( 19 | torch.randn(feature_size, cluster_size) * 1 / math.sqrt(feature_size)) 20 | 21 | self.cluster_weights2 = nn.Parameter( 22 | torch.randn(1, feature_size, cluster_size) * 1 / math.sqrt(feature_size)) 23 | 24 | self.hidden1_weights = nn.Parameter( 25 | torch.randn(feature_size*cluster_size, output_dim)* 1 / math.sqrt(feature_size)) 26 | 27 | if add_batch_norm: 28 | self.cluster_biases = None 29 | self.bn1 = nn.BatchNorm1d(cluster_size) 30 | else: 31 | self.cluster_biases = nn.Parameter(torch.randn(cluster_size) * 1 / math.sqrt(feature_size)) # attention initialization 32 | self.bn1 = None 33 | 34 | self.bn2 = nn.BatchNorm1d(output_dim) 35 | 36 | if gating: 37 | self.context_gating = GatingContext(output_dim, add_batch_norm=add_batch_norm) 38 | 39 | def forward(self, x): 40 | x = x.transpose(1, 3).contiguous() # B x 1024 x N x 1 -> B x 1 x N x 1024 41 | x = x.view((-1, self.max_samples, self.feature_size)) # B x N x 1024 42 | 43 | activation = torch.matmul(x, self.cluster_weights) # B x N x 1024 X 1024 x 64 -> B x N x 64 44 | if self.add_batch_norm: 45 | # activation = activation.transpose(1,2).contiguous() 46 | activation = activation.view(-1, self.cluster_size) # B x N x 64 -> BN x 64 47 | activation = self.bn1(activation) # BN x 64 48 | activation = activation.view(-1, self.max_samples, self.cluster_size) # BN x 64 -> B x N x 64 49 | # activation = activation.transpose(1,2).contiguous() 50 | else: 51 | activation = activation + self.cluster_biases # B x N x 64 + 64 -> B x N x 64 52 | 53 | activation = self.softmax(activation) # B x N x 64 --(dim=-1)--> B x N x 64 54 | # activation = activation[:,:,:64] 55 | activation = activation.view((-1, self.max_samples, self.cluster_size)) # B x N x 64 56 | 57 | a_sum = activation.sum(-2, keepdim=True) # B x N x K --(dim=-2)--> B x 1 x K 58 | a = a_sum * self.cluster_weights2 # B x 1 x K X 1 x C x K -> B x C x K 59 | # element-wise multiply, broadcast mechanism 60 | 61 | 62 | activation = torch.transpose(activation, 2, 1) # B x N x 64 -> B x 64 x N 63 | 64 | x = x.view((-1, self.max_samples, self.feature_size)) # B x N x C -> B x N x C 65 | vlad = torch.matmul(activation, x) # B x K x N X B x N x C -> B x K x C 66 | vlad = torch.transpose(vlad, 2, 1) # B x K x C -> B x C x K 67 | vlad = vlad - a # B x C x K - B x C x K -> B x C x K 68 | 69 | vlad = F.normalize(vlad, dim=1, p=2).contiguous() # B x C x K -> B x C x K 70 | vlad = vlad.view((-1, self.cluster_size * self.feature_size)) # B x (C*K) 71 | return vlad 72 | 73 | class SpatialPyramidNetVLAD(nn.Module): 74 | def __init__(self, feature_size, max_samples, cluster_size, output_dim, 75 | gating=True, add_batch_norm=True): 76 | super(SpatialPyramidNetVLAD, self).__init__() 77 | # max_samples[0] = 64 78 | self.vlad0 = NetVLADBase(feature_size[0], max_samples[0], cluster_size[0], output_dim[0], gating, add_batch_norm) 79 | # max_samples[1] = 256 80 | self.vlad1 = NetVLADBase(feature_size[1], max_samples[1], cluster_size[1], output_dim[1], gating, add_batch_norm) 81 | # max_samples[2] = 1024 82 | self.vlad2 = NetVLADBase(feature_size[2], max_samples[2], cluster_size[2], output_dim[2], gating, add_batch_norm) 83 | # max_samples[3] = 4096 84 | self.vlad3 = NetVLADBase(feature_size[3], max_samples[3], cluster_size[3], output_dim[3], gating, add_batch_norm) 85 | 86 | sum_cluster_size = cluster_size[0] + cluster_size[1] + cluster_size[2] + cluster_size[3] 87 | self.hidden_weights = nn.Parameter(torch.randn(feature_size[0]*sum_cluster_size, output_dim[0])* 1 / math.sqrt(feature_size[0])) 88 | 89 | self.bn2 = nn.BatchNorm1d(output_dim[0]) 90 | self.gating = gating 91 | if self.gating: 92 | self.context_gating = GatingContext(output_dim[0], add_batch_norm=add_batch_norm) 93 | 94 | def forward(self, f0, f1, f2, f3): 95 | v0 = self.vlad0(f0) 96 | v1 = self.vlad1(f1) 97 | v2 = self.vlad2(f2) 98 | v3 = self.vlad3(f3) 99 | vlad = torch.cat((v0, v1, v2, v3), dim=-1) 100 | vlad = torch.matmul(vlad, self.hidden_weights) # B x (1024*64) X (1024*64) x 256 -> B x 256 101 | vlad = self.bn2(vlad) # B x 256 -> B x 256 102 | 103 | if self.gating: 104 | vlad = self.context_gating(vlad) # B x 256 -> B x 256 105 | return vlad # B x 256 106 | 107 | class GatingContext(nn.Module): 108 | def __init__(self, dim, add_batch_norm=True): 109 | super(GatingContext, self).__init__() 110 | self.dim = dim 111 | self.add_batch_norm = add_batch_norm 112 | self.gating_weights = nn.Parameter( 113 | torch.randn(dim, dim) * 1 / math.sqrt(dim)) 114 | self.sigmoid = nn.Sigmoid() 115 | 116 | if add_batch_norm: 117 | self.gating_biases = None 118 | self.bn1 = nn.BatchNorm1d(dim) 119 | else: 120 | self.gating_biases = nn.Parameter( 121 | torch.randn(dim) * 1 / math.sqrt(dim)) 122 | self.bn1 = None 123 | 124 | def forward(self, x): 125 | gates = torch.matmul(x, self.gating_weights) # B x 256 X 256 x 256 -> B x 256 126 | 127 | if self.add_batch_norm: 128 | gates = self.bn1(gates) # B x 256 -> B x 256 129 | else: 130 | gates = gates + self.gating_biases # B x 256 + 256 -> B x 256 131 | 132 | gates = self.sigmoid(gates) # B x 256 -> B x 256 133 | 134 | activation = x * gates # B x 256 * B x 256 -> B x 256 135 | 136 | return activation 137 | -------------------------------------------------------------------------------- /models/model_factory.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os, sys 3 | sys.path.append(os.path.join(os.path.dirname(__file__), '../')) 4 | # from model import pointnet_cls, pointnet2_cls_msg, DetectAndVLAD,lightPNSA, lightPN,lightPNplus, pyramid, RandLANet 5 | 6 | # from model import DetectAndVLAD, NetMG2, Backbone 7 | 8 | from models import pptnet_v2 9 | 10 | from models import PointNetVlad as PNV 11 | 12 | from models.MinkLoc3dv2.layers.eca_block import ECABasicBlock 13 | from models.MinkLoc3dv2.minkfpn import MinkFPN 14 | from models.MinkLoc3dv2.layers.pooling_wrapper import PoolingWrapper 15 | from models.MinkLoc3dv2.mink_params import create_resnet_block 16 | from models.MinkLoc3dv2.minkloc import MinkLoc, MinkLocLAWS 17 | 18 | def model_factory(args): 19 | #### Model 20 | print(args.backbone) 21 | if args.backbone == 'ppt': 22 | params = pptnet_v2.PPTparams() 23 | net = pptnet_v2.Network(param=params) 24 | if args.backbone == 'ppt_laws': 25 | params = pptnet_v2.PPTparams() 26 | net = pptnet_v2.NetworkLAWS(param=params) 27 | 28 | if args.backbone == 'pnv': 29 | net = PNV.PointNetVlad(num_points=4096, 30 | global_feat=True, 31 | feature_transform=True, 32 | max_pool=False, 33 | output_dim=256, 34 | use_rgb = args.use_rgb) 35 | if args.backbone == 'pnv_laws': 36 | net = PNV.PointNetVladLAWScopy(num_points=4096, 37 | global_feat=True, 38 | feature_transform=True, 39 | max_pool=False, 40 | output_dim=256, 41 | use_rgb = args.use_rgb) 42 | if args.backbone == 'mink': 43 | in_channels = args.input_channel 44 | block_module = create_resnet_block(args.block) 45 | backbone = MinkFPN(in_channels=in_channels, out_channels=args.feature_size, 46 | num_top_down=args.num_top_down, conv0_kernel_size=args.conv0_kernel_size, 47 | block=block_module, layers=args.layers, planes=args.planes) 48 | pooling = PoolingWrapper(pool_method=args.pooling, in_dim=args.feature_size, 49 | output_dim=args.output_dim) 50 | net = MinkLoc(backbone=backbone, pooling=pooling, normalize_embeddings=args.normalize_embeddings) 51 | if args.backbone == 'mink_laws': 52 | in_channels = args.input_channel 53 | block_module = create_resnet_block(args.block) 54 | backbone = MinkFPN(in_channels=in_channels, out_channels=args.feature_size, 55 | num_top_down=args.num_top_down, conv0_kernel_size=args.conv0_kernel_size, 56 | block=block_module, layers=args.layers, planes=args.planes) 57 | net = MinkLocLAWS(backbone=backbone, 58 | pool_method = args.pooling, 59 | in_dim=args.feature_size, 60 | output_dim=args.output_dim, 61 | normalize_embeddings=args.normalize_embeddings) 62 | 63 | return net 64 | 65 | -------------------------------------------------------------------------------- /train.sh: -------------------------------------------------------------------------------- 1 | # #!/bin/sh 2 | # export PYTHONPATH=./ 3 | 4 | 5 | # exp_name=ppt_mulran_cosplace 6 | # exp_dir=exp/${exp_name} 7 | # mkdir -p ${exp_dir} 8 | 9 | # trainer: BasicTrainer BackboneExpansion 10 | # dataset: Oxforf/Muran: TrainDataset TrainDatasetv2 TrainDatasetv0 | ScannetPR: ScannetPRDataset, ScannetPRDatasetv2 --use_rgb \ 11 | # backbone:ppt ppt_laws --input_channel 3 \ | pnv pnv_laws --input_channel 3 \ | mink_laws --input_channel 1 \ 12 | # dataset_folder /home/xy/xy/code/Oxford/data/benchmark_datasets | /home/xy/xy/code/Mulran | /home/xy/xy/code/Data/ScannetPR | 13 | 14 | python train.py \ 15 | --trainer BackboneExpansion \ 16 | --dataset_folder /home/xy/xy/code/Oxford/data/benchmark_datasets \ 17 | --dataset TrainDatasetv2 \ 18 | --backbone pnv_laws \ 19 | --save_dir pnv_test/ \ 20 | --pooling netvlad \ 21 | --normalize_embeddings \ 22 | --input_channel 3 \ 23 | --batch_size 16 \ 24 | --iterations_per_epoch 10 \ 25 | --epochs_num 16 \ 26 | --fc_output_dim 256 \ 27 | --lr 0.001 \ 28 | --M 20 \ 29 | --N 2 \ 30 | --groups_num 8 \ 31 | --min_images_per_class 2 \ 32 | 33 | -------------------------------------------------------------------------------- /util/__pycache__/computer_overlap.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/computer_overlap.cpython-38.pyc -------------------------------------------------------------------------------- /util/__pycache__/loading_pointclouds.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/loading_pointclouds.cpython-36.pyc -------------------------------------------------------------------------------- /util/__pycache__/loading_pointclouds.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/loading_pointclouds.cpython-38.pyc -------------------------------------------------------------------------------- /util/__pycache__/parser.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/parser.cpython-36.pyc -------------------------------------------------------------------------------- /util/__pycache__/parser.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/parser.cpython-38.pyc -------------------------------------------------------------------------------- /util/__pycache__/pointnetvlad_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/pointnetvlad_loss.cpython-36.pyc -------------------------------------------------------------------------------- /util/__pycache__/pointnetvlad_loss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/pointnetvlad_loss.cpython-38.pyc -------------------------------------------------------------------------------- /util/__pycache__/pt_util.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/pt_util.cpython-36.pyc -------------------------------------------------------------------------------- /util/__pycache__/pt_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/pt_util.cpython-38.pyc -------------------------------------------------------------------------------- /util/__pycache__/scannet_config.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/scannet_config.cpython-38.pyc -------------------------------------------------------------------------------- /util/__pycache__/util.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/util.cpython-36.pyc -------------------------------------------------------------------------------- /util/__pycache__/util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WHU-USI3DV/LAWS/f65b6237804658d99af16c2cc1ad132bef9a1631/util/__pycache__/util.cpython-38.pyc -------------------------------------------------------------------------------- /util/util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import time 4 | import torch 5 | import random 6 | 7 | class AverageMeter(object): 8 | """Computes and stores the average and current value""" 9 | def __init__(self): 10 | self.reset() 11 | 12 | def reset(self): 13 | self.val = 0 14 | self.avg = 0 15 | self.sum = 0 16 | self.count = 0 17 | 18 | def update(self, val, n=1): 19 | self.val = val 20 | self.sum += val * n 21 | self.count += n 22 | self.avg = self.sum / self.count 23 | 24 | def check_mkdir(dir_name): 25 | if not os.path.exists(dir_name): 26 | os.mkdir(dir_name) 27 | 28 | 29 | def check_makedirs(dir_name): 30 | if not os.path.exists(dir_name): 31 | os.makedirs(dir_name) 32 | 33 | def count_files(): 34 | root = '/test/dataset/benchmark_datasets/oxford' # dataset path 35 | files = os.listdir(root) 36 | cnt = 0 37 | for i in range(len(files)): 38 | data_path = os.path.join(root, files[i], 'pointcloud_20m_10overlap') 39 | cnt += len(os.listdir(data_path)) 40 | print('data files: {}'.format(cnt)) 41 | return cnt 42 | 43 | def plot_point_cloud(points, label=None, output_filename=''): 44 | """ points is a Nx3 numpy array """ 45 | 46 | import matplotlib.pyplot as plt 47 | # from mpl_toolkits.mplot3d import Axes3D 48 | fig = plt.figure() 49 | # ax = Axes3D(fig) 50 | ax = plt.subplot(111, projection='3d') 51 | if label is not None: 52 | point = ax.scatter(points[:, 0], points[:, 1], points[:, 2], 53 | # cmap='RdYlBu', 54 | c = label, 55 | # linewidth=2, 56 | alpha=0.5, 57 | marker=".") 58 | else: 59 | point = ax.scatter(points[:, 0], points[:, 1], points[:, 2], 60 | # cmap='RdYlBu', 61 | c = points[:, 2], 62 | # linewidth=2, 63 | alpha=0.5, 64 | marker=".") 65 | ax.set_xlabel('x') 66 | ax.set_ylabel('y') 67 | ax.set_zlabel('z') 68 | # fig.colorbar(point) 69 | # plt.axis('scaled') 70 | # plt.axis('off') 71 | if output_filename!='': 72 | plt.savefig(output_filename, dpi=300, bbox_inches='tight') 73 | plt.close() 74 | else: 75 | plt.show() 76 | 77 | class Timer(object): 78 | """A simple timer.""" 79 | # Ref: https://github.com/chrischoy/FCGF/blob/master/lib/timer.py 80 | 81 | def __init__(self, binary_fn=None, init_val=0): 82 | self.total_time = 0. 83 | self.calls = 0 84 | self.start_time = 0. 85 | self.diff = 0. 86 | self.binary_fn = binary_fn 87 | self.tmp = init_val 88 | 89 | def reset(self): 90 | self.total_time = 0 91 | self.calls = 0 92 | self.start_time = 0 93 | self.diff = 0 94 | 95 | @property 96 | def avg(self): 97 | return self.total_time / self.calls 98 | 99 | def tic(self): 100 | # using time.time instead of time.clock because time time.clock 101 | # does not normalize for multithreading 102 | self.start_time = time.time() 103 | 104 | def toc(self, average=True): 105 | self.diff = time.time() - self.start_time 106 | self.total_time += self.diff 107 | self.calls += 1 108 | if self.binary_fn: 109 | self.tmp = self.binary_fn(self.tmp, self.diff) 110 | if average: 111 | return self.avg 112 | else: 113 | return self.diff 114 | 115 | 116 | def make_deterministic(seed=0): 117 | """Make results deterministic. If seed == -1, do not make deterministic. 118 | Running your script in a deterministic way might slow it down. 119 | Note that for some packages (eg: sklearn's PCA) this function is not enough. 120 | """ 121 | seed = int(seed) 122 | if seed == -1: 123 | return 124 | random.seed(seed) 125 | np.random.seed(seed) 126 | torch.manual_seed(seed) 127 | torch.cuda.manual_seed_all(seed) 128 | torch.backends.cudnn.deterministic = True 129 | torch.backends.cudnn.benchmark = False 130 | 131 | 132 | def setup_logging(output_folder, exist_ok=False, console="debug", 133 | info_filename="info.log", debug_filename="debug.log"): 134 | """Set up logging files and console output. 135 | Creates one file for INFO logs and one for DEBUG logs. 136 | Args: 137 | output_folder (str): creates the folder where to save the files. 138 | exist_ok (boolean): if False throw a FileExistsError if output_folder already exists 139 | debug (str): 140 | if == "debug" prints on console debug messages and higher 141 | if == "info" prints on console info messages and higher 142 | if == None does not use console (useful when a logger has already been set) 143 | info_filename (str): the name of the info file. if None, don't create info file 144 | debug_filename (str): the name of the debug file. if None, don't create debug file 145 | """ 146 | import os 147 | import sys 148 | import logging 149 | import traceback 150 | if not exist_ok and os.path.exists(output_folder): 151 | raise FileExistsError(f"{output_folder} already exists!") 152 | os.makedirs(output_folder, exist_ok=True) 153 | base_formatter = logging.Formatter('%(asctime)s %(message)s', "%Y-%m-%d %H:%M:%S") 154 | logger = logging.getLogger('') 155 | logger.setLevel(logging.DEBUG) 156 | 157 | if info_filename != None: 158 | info_file_handler = logging.FileHandler(f'{output_folder}/{info_filename}') 159 | info_file_handler.setLevel(logging.INFO) 160 | info_file_handler.setFormatter(base_formatter) 161 | logger.addHandler(info_file_handler) 162 | 163 | if debug_filename != None: 164 | debug_file_handler = logging.FileHandler(f'{output_folder}/{debug_filename}') 165 | debug_file_handler.setLevel(logging.DEBUG) 166 | debug_file_handler.setFormatter(base_formatter) 167 | logger.addHandler(debug_file_handler) 168 | 169 | if console != None: 170 | console_handler = logging.StreamHandler() 171 | if console == "debug": console_handler.setLevel(logging.DEBUG) 172 | if console == "info": console_handler.setLevel(logging.INFO) 173 | console_handler.setFormatter(base_formatter) 174 | logger.addHandler(console_handler) 175 | 176 | def my_handler(type_, value, tb): 177 | logger.info("\n" + "".join(traceback.format_exception(type, value, tb))) 178 | logging.info("Experiment finished (with some errors)") 179 | sys.excepthook = my_handler 180 | 181 | 182 | class InfiniteDataLoader(torch.utils.data.DataLoader): 183 | def __init__(self, *args, **kwargs): 184 | super().__init__(*args, **kwargs) 185 | self.dataset_iterator = super().__iter__() 186 | 187 | def __iter__(self): 188 | return self 189 | 190 | def __next__(self): 191 | try: 192 | batch = next(self.dataset_iterator) 193 | except StopIteration: 194 | self.dataset_iterator = super().__iter__() 195 | batch = next(self.dataset_iterator) 196 | return batch 197 | 198 | 199 | class AverageMeter(object): 200 | """Computes and stores the average and current value""" 201 | 202 | def __init__(self): 203 | self.reset() 204 | 205 | def reset(self): 206 | self.val = 0 207 | self.avg = 0 208 | self.sum = 0 209 | self.count = 0 210 | 211 | def update(self, val, n=1): 212 | self.val = val 213 | self.sum += val * n 214 | self.count += n 215 | self.avg = self.sum / self.count 216 | 217 | 218 | def move_to_device(optimizer, device): 219 | for state in optimizer.state.values(): 220 | for k, v in state.items(): 221 | if torch.is_tensor(v): 222 | state[k] = v.to(device) 223 | 224 | 225 | def save_checkpoint(state, is_best, output_folder, ckpt_filename="last_checkpoint.pth"): 226 | # TODO it would be better to move weights to cpu before saving 227 | checkpoint_path = f"{output_folder}/{ckpt_filename}" 228 | torch.save(state, checkpoint_path) 229 | if is_best: 230 | torch.save(state["model_state_dict"], f"{output_folder}/best_model.pth") 231 | --------------------------------------------------------------------------------