diff --git a/faiss/gpu/GpuIndexIVF.cu b/faiss/gpu/GpuIndexIVF.cu index a549703c21..96be02589a 100644 --- a/faiss/gpu/GpuIndexIVF.cu +++ b/faiss/gpu/GpuIndexIVF.cu @@ -207,6 +207,8 @@ void GpuIndexIVF::copyFrom(const faiss::IndexIVF* index) { "that already contains a GPU coarse (level 1) quantizer " "is not currently supported"); } + // copy the num vectors to use to train the encoder from the source index + encoderTrainingVecs_ = index->train_encoder_num_vectors(); // Validate equality FAISS_ASSERT(is_trained == index->is_trained); @@ -254,6 +256,10 @@ void GpuIndexIVF::copyTo(faiss::IndexIVF* index) const { index->make_direct_map(false); } +idx_t GpuIndexIVF::train_encoder_num_vectors() const { + return encoderTrainingVecs_; +} + idx_t GpuIndexIVF::getNumLists() const { return nlist; } diff --git a/faiss/gpu/GpuIndexIVF.h b/faiss/gpu/GpuIndexIVF.h index d6fd5b6ffa..b593a0a57f 100644 --- a/faiss/gpu/GpuIndexIVF.h +++ b/faiss/gpu/GpuIndexIVF.h @@ -81,6 +81,10 @@ class GpuIndexIVF : public GpuIndex, public IndexIVFInterface { /// Returns the number of inverted lists we're managing virtual idx_t getNumLists() const; + /// can be redefined by subclasses to indicate how many training vectors + /// they need + virtual idx_t train_encoder_num_vectors() const; + /// Returns the number of vectors present in a particular inverted list virtual idx_t getListLength(idx_t listId) const; @@ -148,6 +152,9 @@ class GpuIndexIVF : public GpuIndex, public IndexIVFInterface { /// For a trained/initialized index, this is a reference to the base class std::shared_ptr baseIndex_; + + /// Number of vectors to use to train the encoder (if applicable) + idx_t encoderTrainingVecs_ = 0; }; } // namespace gpu diff --git a/faiss/gpu/GpuIndexIVFScalarQuantizer.cu b/faiss/gpu/GpuIndexIVFScalarQuantizer.cu index 4d4c3c7b4c..551c90eee8 100644 --- a/faiss/gpu/GpuIndexIVFScalarQuantizer.cu +++ b/faiss/gpu/GpuIndexIVFScalarQuantizer.cu @@ -12,6 +12,7 @@ #include #include #include +#include #include namespace faiss { @@ -219,14 +220,19 @@ void GpuIndexIVFScalarQuantizer::reset() { void GpuIndexIVFScalarQuantizer::trainResiduals_(idx_t n, const float* x) { // The input is already guaranteed to be on the CPU + idx_t max_nt = train_encoder_num_vectors(); + if (max_nt <= 0) { + max_nt = (size_t)1 << 35; + } + TransformedVectors tv(x, fvecs_maybe_subsample(d, (size_t*)&n, max_nt, x, verbose)); if (!by_residual) { - sq.train(n, x); + sq.train(n, tv.x); } else { std::vector assign(n); - quantizer->assign(n, x, assign.data()); + quantizer->assign(n, tv.x, assign.data()); std::vector residuals(n * d); - quantizer->compute_residual_n(n, x, residuals.data(), assign.data()); + quantizer->compute_residual_n(n, tv.x, residuals.data(), assign.data()); sq.train(n, residuals.data()); }