| 74 | self.dist_func = dist_type_to_dist_func[self.knn_sim_func] # l2 or dot product function |
| 75 | |
| 76 | def setup_faiss(self): |
| 77 | if not self.val_file: |
| 78 | raise ValueError('Cannot build a datastore without the data.') |
| 79 | |
| 80 | start = time.time() |
| 81 | cpu_index = faiss.read_index(self.index_file, faiss.IO_FLAG_ONDISK_SAME_DIR) |
| 82 | logger.info(f'Reading datastore took {time.time() - start} s') |
| 83 | cpu_index.nprobe = self.probe |
| 84 | |
| 85 | if self.knn_gpu: |
| 86 | start = time.time() |
| 87 | # 1. Set maximum GPU memory allocation |
| 88 | res = faiss.StandardGpuResources() |
| 89 | res.setTempMemory(40 * 1024 * 1024 * 1024) |
| 90 | |
| 91 | # 2. Create a more memory-efficient cloner |
| 92 | co = faiss.GpuClonerOptions() |
| 93 | co.useFloat16 = True # Use half precision |
| 94 | co.verbose = True # See what's happening |
| 95 | |
| 96 | gpu_index = faiss.index_cpu_to_gpu(res, self.local_process_index, cpu_index, co) |
| 97 | |
| 98 | logger.info(f'Moving index to GPU took {time.time() - start} s') |
| 99 | else: |
| 100 | gpu_index = cpu_index |
| 101 | |
| 102 | # make_direct_map() allows calling reconstruct(n), |
| 103 | # and reconstructing key vectors given their ids |
| 104 | # currently, this is implemented only for CPU indexes: |
| 105 | # https://github.com/facebookresearch/faiss/issues/2181 |
| 106 | cpu_index.make_direct_map() |
| 107 | |
| 108 | start_time = time.time() |
| 109 | # Load val_file using pickle |
| 110 | with open(self.val_file, 'rb') as f: |
| 111 | self.vals = pickle.load(f).to(self.device) |
| 112 | logger.info(f'Loading keys and vals to memory took {time.time() - start_time} s') |
| 113 | |
| 114 | return cpu_index, gpu_index |
| 115 | |
| 116 | def break_into(self, model): |
| 117 | self.model = model |