| 132 | } |
| 133 | |
| 134 | static std::vector<size_t> get_nprobes( |
| 135 | const index_type& ivf, |
| 136 | const std::vector<size_t>& all_nprobes, |
| 137 | data_type& query, |
| 138 | gt_type& gt |
| 139 | ) { |
| 140 | size_t nq = query.rows(); |
| 141 | size_t total_count = topk * nq; |
| 142 | float old_recall = 0; |
| 143 | std::vector<size_t> nprobes; |
| 144 | |
| 145 | for (auto nprobe : all_nprobes) { |
| 146 | nprobes.push_back(nprobe); |
| 147 | |
| 148 | size_t total_correct = 0; |
| 149 | std::vector<PID> results(topk); |
| 150 | for (size_t i = 0; i < nq; i++) { |
| 151 | ivf.search(&query(i, 0), topk, nprobe, results.data()); |
| 152 | for (size_t j = 0; j < topk; j++) { |
| 153 | for (size_t k = 0; k < topk; k++) { |
| 154 | if (gt(i, k) == results[j]) { |
| 155 | total_correct++; |
| 156 | break; |
| 157 | } |
| 158 | } |
| 159 | } |
| 160 | } |
| 161 | float recall = static_cast<float>(total_correct) / static_cast<float>(total_count); |
| 162 | if (recall > 0.997 || recall - old_recall < 1e-5) { |
| 163 | break; |
| 164 | } |
| 165 | std::cout << recall << '\t' << nprobe << std::endl << std::flush; |
| 166 | old_recall = recall; |
| 167 | } |
| 168 | |
| 169 | return nprobes; |
| 170 | } |