* @brief Predict logits for samples * @param _samples ndarray of edge samples, with shape (?, kSampleSize) */
| 727 | * @param _samples ndarray of edge samples, with shape (?, kSampleSize) |
| 728 | */ |
| 729 | py::array_t<Float> predict_numpy(const py::array_t<Index> &_array) { |
| 730 | if (_array.ndim() != 2 || _array.shape(1) != kSampleSize) { |
| 731 | std::stringstream ss; |
| 732 | ss << _array.shape(0); |
| 733 | for (int i = 1; i < _array.ndim(); i++) |
| 734 | ss << ", " << _array.shape(i); |
| 735 | LOG(FATAL) << "Expect an array with shape (?, " << kSampleSize |
| 736 | << "), but shape (" << ss.str() << ") is found"; |
| 737 | } |
| 738 | array = &_array; |
| 739 | is_train = false; |
| 740 | |
| 741 | pool_offsets.resize(num_sampler + 1); |
| 742 | for (auto &&sampler_offsets : pool_offsets) { |
| 743 | sampler_offsets.resize(num_partition); |
| 744 | for (auto &&partition_offsets : sampler_offsets) |
| 745 | partition_offsets.resize(num_partition); |
| 746 | } |
| 747 | predict_pool.resize(num_partition); |
| 748 | for (auto &&partition_pool : predict_pool) |
| 749 | partition_pool.resize(num_partition); |
| 750 | sample_indexes.resize(num_partition); |
| 751 | for (auto &&partition_indexes : sample_indexes) |
| 752 | partition_indexes.resize(num_partition); |
| 753 | |
| 754 | std::vector<std::thread> sample_threads(num_sampler + num_worker); |
| 755 | size_t num_sample = array->shape(0); |
| 756 | size_t work_load = (num_sample + num_sampler - 1) / num_sampler; |
| 757 | for (int i = 0; i < num_sampler + num_worker; i++) |
| 758 | sample_threads[i] = std::thread(&Sampler::count_numpy, samplers[0], work_load * i, |
| 759 | std::min(work_load * (i + 1), num_sample), i); |
| 760 | for (auto &&thread : sample_threads) |
| 761 | thread.join(); |
| 762 | |
| 763 | for (int i = 0; i < num_sampler; i++) |
| 764 | for (int j = 0; j < num_partition; j++) |
| 765 | for (int k = 0; k < num_partition; k++) |
| 766 | pool_offsets[i + 1][j][k] += pool_offsets[i][j][k]; |
| 767 | predict_batch_id = 0; |
| 768 | num_predict_batch = 0; |
| 769 | size_t all_pool = 0; |
| 770 | for (int i = 0; i < num_partition; i++) |
| 771 | for (int j = 0; j < num_partition; j++) { |
| 772 | size_t this_pool_size = pool_offsets[num_sampler][i][j]; |
| 773 | all_pool += this_pool_size; |
| 774 | predict_pool[i][j].resize(this_pool_size); |
| 775 | sample_indexes[i][j].resize(this_pool_size); |
| 776 | num_predict_batch += (this_pool_size + batch_size - 1) / batch_size; |
| 777 | } |
| 778 | |
| 779 | for (int i = 0; i < num_sampler + num_worker; i++) |
| 780 | sample_threads[i] = std::thread(&Sampler::distribute_numpy, samplers[0], work_load * i, |
| 781 | std::min(work_load * (i + 1), num_sample), i); |
| 782 | for (auto &&thread : sample_threads) |
| 783 | thread.join(); |
| 784 | |
| 785 | results.resize(num_sample); |
| 786 | std::vector<std::thread> worker_threads(num_worker); |