Python faiss.METRIC_L2 Examples

def _build_approximate_index(self,
                                     data: np.ndarray):
            dimensionality = data.shape[1]
            nlist = 100 if data.shape[0] > 100 else 2

            if self.kernel_name in {'rbf'}:
                quantizer = faiss.IndexFlatL2(dimensionality)
                cpu_index_flat = faiss.IndexIVFFlat(quantizer, dimensionality, nlist, faiss.METRIC_L2)
            else:
                quantizer = faiss.IndexFlatIP(dimensionality)
                cpu_index_flat = faiss.IndexIVFFlat(quantizer, dimensionality, nlist)

            gpu_index_ivf = faiss.index_cpu_to_gpu(self.resource, 0, cpu_index_flat)
            gpu_index_ivf.train(data)
            gpu_index_ivf.add(data)
            self.index = gpu_index_ivf 

def train_index(data, quantizer_path, trained_index_path, fine_quant='SQ8', cuda=False):
    quantizer = faiss.read_index(quantizer_path)
    if fine_quant == 'SQ8':
        trained_index = faiss.IndexIVFScalarQuantizer(quantizer, quantizer.d, quantizer.ntotal, faiss.METRIC_L2)
    elif fine_quant.startswith('PQ'):
        m = int(fine_quant[2:])
        trained_index = faiss.IndexIVFPQ(quantizer, quantizer.d, quantizer.ntotal, m, 8)
    else:
        raise ValueError(fine_quant)

    if cuda:
        if fine_quant.startswith('PQ'):
            print('PQ not supported on GPU; keeping CPU.')
        else:
            res = faiss.StandardGpuResources()
            gpu_index = faiss.index_cpu_to_gpu(res, 0, trained_index)
            gpu_index.train(data)
            trained_index = faiss.index_gpu_to_cpu(gpu_index)
    else:
        trained_index.train(data)
    faiss.write_index(trained_index, trained_index_path) 

def __init__(self, cell_size=20, nr_cells=1024, K=4, num_lists=32, probes=32, res=None, train=None, gpu_id=-1):
    super(FAISSIndex, self).__init__()
    self.cell_size = cell_size
    self.nr_cells = nr_cells
    self.probes = probes
    self.K = K
    self.num_lists = num_lists
    self.gpu_id = gpu_id

    # BEWARE: if this variable gets deallocated, FAISS crashes
    self.res = res if res else faiss.StandardGpuResources()
    self.res.setTempMemoryFraction(0.01)
    if self.gpu_id != -1:
      self.res.initializeForDevice(self.gpu_id)

    nr_samples = self.nr_cells * 100 * self.cell_size
    train = train if train is not None else T.randn(self.nr_cells * 100, self.cell_size)

    self.index = faiss.GpuIndexIVFFlat(self.res, self.cell_size, self.num_lists, faiss.METRIC_L2)
    self.index.setNumProbes(self.probes)
    self.train(train) 

def fit(self, X):
        if self._metric == 'angular':
            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')

        if X.dtype != numpy.float32:
            X = X.astype(numpy.float32)

        self.quantizer = faiss.IndexFlatL2(X.shape[1])
        index = faiss.IndexIVFFlat(
            self.quantizer, X.shape[1], self._n_list, faiss.METRIC_L2)
        index.train(X)
        index.add(X)
        self.index = index 

def fit(self, X):
        X = X.astype(numpy.float32)
        self._index = faiss.GpuIndexIVFFlat(self._res, len(X[0]), self._n_bits,
                                            faiss.METRIC_L2)
        # self._index = faiss.index_factory(len(X[0]),
        #                                   "IVF%d,Flat" % self._n_bits)
        # co = faiss.GpuClonerOptions()
        # co.useFloat16 = True
        # self._index = faiss.index_cpu_to_gpu(self._res, 0,
        #                                      self._index, co)
        self._index.train(X)
        self._index.add(X)
        self._index.setNumProbes(self._n_probes) 

def _faiss_knn(keys: torch.Tensor,
               queries: torch.Tensor,
               num_neighbors: int,
               distance: str) -> Tuple[torch.Tensor, torch.Tensor]:
    # https://github.com/facebookresearch/XLM/blob/master/src/model/memory/utils.py
    if not is_faiss_available():
        raise RuntimeError("faiss_knn requires faiss-gpu")
    import faiss

    assert distance in ['dot_product', 'l2']
    assert keys.size(1) == queries.size(1)

    metric = faiss.METRIC_INNER_PRODUCT if distance == 'dot_product' else faiss.METRIC_L2

    k_ptr = _tensor_to_ptr(keys)
    q_ptr = _tensor_to_ptr(queries)

    scores = keys.new_zeros((queries.size(0), num_neighbors), dtype=torch.float32)
    indices = keys.new_zeros((queries.size(0), num_neighbors), dtype=torch.int64)

    s_ptr = _tensor_to_ptr(scores)
    i_ptr = _tensor_to_ptr(indices)

    faiss.bruteForceKnn(FAISS_RES, metric,
                        k_ptr, True, keys.size(0),
                        q_ptr, True, queries.size(0),
                        queries.size(1), num_neighbors, s_ptr, i_ptr)
    return scores, indices 

def execute(cls, ctx, op):
        (y,), device_id, xp = as_same_device(
            [ctx[op.input.key]], device=op.device, ret_extra=True)
        indexes = [_load_index(ctx, op, ctx[index.key], device_id)
                   for index in op.inputs[1:]]

        with device(device_id):
            y = xp.ascontiguousarray(y, dtype=np.float32)

            if len(indexes) == 1:
                index = indexes[0]
            else:
                index = faiss.IndexShards(indexes[0].d)
                [index.add_shard(ind) for ind in indexes]

            if op.metric == 'cosine':
                # faiss does not support cosine distances directly,
                # data needs to be normalize before searching,
                # refer to:
                # https://github.com/facebookresearch/faiss/wiki/FAQ#how-can-i-index-vectors-for-cosine-distance
                faiss.normalize_L2(y)

            if op.nprobe is not None:
                index.nprobe = op.nprobe

            if device_id >= 0:  # pragma: no cover
                n = y.shape[0]
                k = op.n_neighbors
                distances = xp.empty((n, k), dtype=xp.float32)
                indices = xp.empty((n, k), dtype=xp.int64)
                index.search_c(n, _swig_ptr_from_cupy_float32_array(y),
                               k, _swig_ptr_from_cupy_float32_array(distances),
                               _swig_ptr_from_cupy_int64_array(indices))
            else:
                distances, indices = index.search(y, op.n_neighbors)
            if op.return_distance:
                if index.metric_type == faiss.METRIC_L2:
                    # make it equivalent to `pairwise.euclidean_distances`
                    distances = xp.sqrt(distances, out=distances)
                elif op.metric == 'cosine':
                    # make it equivalent to `pairwise.cosine_distances`
                    distances = xp.subtract(1, distances, out=distances)
                ctx[op.outputs[0].key] = distances
            ctx[op.outputs[-1].key] = indices