faiss_product_quantization.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOTWlOX8UHfyY3QdhalLe9l",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/weedge/doraemon-nb/blob/main/faiss_product_quantization.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "https://www.pinecone.io/learn/series/faiss/product-quantization/\n",
        "\n",
        "https://www.youtube.com/watch?v=t9mRf2S5vDI\n",
        "https://www.youtube.com/watch?v=BMYBwbkbVec\n"
      ],
      "metadata": {
        "id": "25PljjlaKqmg"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EbyPLVonKkB6",
        "outputId": "b2a265fb-2eea-411e-fcf3-79b29e6b71e8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree... Done\n",
            "Reading state information... Done\n",
            "The following additional packages will be installed:\n",
            "  libomp-14-dev libomp5-14\n",
            "Suggested packages:\n",
            "  libomp-14-doc\n",
            "The following NEW packages will be installed:\n",
            "  libomp-14-dev libomp-dev libomp5-14\n",
            "0 upgraded, 3 newly installed, 0 to remove and 18 not upgraded.\n",
            "Need to get 738 kB of archives.\n",
            "After this operation, 8,991 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu jammy-updates/universe amd64 libomp5-14 amd64 1:14.0.0-1ubuntu1.1 [389 kB]\n",
            "Get:2 http://archive.ubuntu.com/ubuntu jammy-updates/universe amd64 libomp-14-dev amd64 1:14.0.0-1ubuntu1.1 [347 kB]\n",
            "Get:3 http://archive.ubuntu.com/ubuntu jammy/universe amd64 libomp-dev amd64 1:14.0-55~exp2 [3,074 B]\n",
            "Fetched 738 kB in 1s (1,129 kB/s)\n",
            "Selecting previously unselected package libomp5-14:amd64.\n",
            "(Reading database ... 120895 files and directories currently installed.)\n",
            "Preparing to unpack .../libomp5-14_1%3a14.0.0-1ubuntu1.1_amd64.deb ...\n",
            "Unpacking libomp5-14:amd64 (1:14.0.0-1ubuntu1.1) ...\n",
            "Selecting previously unselected package libomp-14-dev.\n",
            "Preparing to unpack .../libomp-14-dev_1%3a14.0.0-1ubuntu1.1_amd64.deb ...\n",
            "Unpacking libomp-14-dev (1:14.0.0-1ubuntu1.1) ...\n",
            "Selecting previously unselected package libomp-dev:amd64.\n",
            "Preparing to unpack .../libomp-dev_1%3a14.0-55~exp2_amd64.deb ...\n",
            "Unpacking libomp-dev:amd64 (1:14.0-55~exp2) ...\n",
            "Setting up libomp5-14:amd64 (1:14.0.0-1ubuntu1.1) ...\n",
            "Setting up libomp-14-dev (1:14.0.0-1ubuntu1.1) ...\n",
            "Setting up libomp-dev:amd64 (1:14.0-55~exp2) ...\n",
            "Processing triggers for libc-bin (2.35-0ubuntu3.1) ...\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_0.so.3 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbb.so.12 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc_proxy.so.2 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc.so.2 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind.so.3 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_5.so.3 is not a symbolic link\n",
            "\n",
            "Collecting faiss-cpu\n",
            "  Downloading faiss_cpu-1.7.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (17.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m17.6/17.6 MB\u001b[0m \u001b[31m49.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting faiss-gpu\n",
            "  Downloading faiss_gpu-1.7.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (85.5 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m85.5/85.5 MB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: faiss-gpu, faiss-cpu\n",
            "Successfully installed faiss-cpu-1.7.4 faiss-gpu-1.7.2\n"
          ]
        }
      ],
      "source": [
        "!apt install libomp-dev\n",
        "!pip install --upgrade faiss-cpu faiss-gpu"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "矢量相似性搜索可能需要大量内存。包含 1M 密集向量（当今世界的小型数据集）的索引通常需要几 GB 内存来存储。\n",
        "\n",
        "高维数据加剧了内存使用过多的问题，并且随着数据集大小的不断增加，这很快就会变得难以管理。\n",
        "\n",
        "乘积量化 (PQ) 是一种流行的方法，可显着压缩高维向量以减少 97% 的内存使用，并且在我们的测试中使最近邻搜索速度加快 5.5 倍。\n",
        "\n",
        "复合 IVF+PQ 索引在不影响准确性的情况下将搜索速度再提高 16.5 倍，与非量化索引相比，总速度提高了 92 倍。\n",
        "\n",
        "在本文中，我们将介绍您需要了解的有关 PQ 的所有内容：它是如何工作的、优缺点、Faiss 中的实现、复合 IVFPQ 索引以及如何实现上述速度和内存优化。\n",
        "\n",
        "## 什么是量化\n",
        "\n",
        "量化是一种通用方法，指将数据压缩到更小的空间中。我知道这可能没有多大意义——让我解释一下。\n",
        "\n",
        "首先，我们来谈谈降维——这与量化*不同*。\n",
        "\n",
        "假设我们有一个高维向量，它的维数为**128**。这些值是 32 位浮点数，范围为*0.0 -> 157.0*（我们的范围**S**）。通过降维，我们的目标是产生另一个更低维度的向量。\n",
        "\n",
        "![降维会降低向量的维数 D，但不会降低范围 S。](https://cdn.sanity.io/images/vr8gru94/production/178d6efe3bfeae550d18bcd8bd3352817717674c-1920x850.png)\n",
        "\n",
        "降维会降低向量的维数 D，但不会降低范围 S。\n",
        "\n",
        "另一方面，我们有量化。量化不关心维数**D**。相反，它针对的是潜在的价值范围。我们不是减少**D**，而是减少**S**。\n",
        "\n",
        "![量化减少了可能向量的范围 S。 请注意，通过预量化，范围通常是无限的。](https://cdn.sanity.io/images/vr8gru94/production/ca04d5d84cd168d0c6a9dba146851ebde6612ead-1920x1080.png)\n",
        "\n",
        "量化减少了可能向量的范围 S。请注意，通过预量化，范围通常是无限的。\n",
        "\n",
        "有很多方法可以做到这一点。例如，我们有*聚类*。当我们对一组向量进行聚类时，我们用较小的*离散和符号质*心集替换更大范围的潜在值（所有可能的向量）。\n",
        "\n",
        "这实际上就是我们定义量化操作的方式。将向量变换为具有有限数量可能值的空间，其中这些值是原始向量的符号表示*。*\n",
        "\n",
        "需要明确的是，这些符号表示的形式各不相同。它们可以是质心（如 PQ 的情况），也可以是[二进制代码（如 LSH 生成的代码）](https://www.pinecone.io/learn/series/faiss/locality-sensitive-hashing-random-projection/)。\n",
        "\n",
        "### 为什么要产品量化？\n",
        "\n",
        "量化主要用于减少索引的内存占用——在比较大型向量数组时这是一项重要任务，因为它们必须全部加载到内存中才能进行比较。\n",
        "\n",
        "然而，PQ 并不是唯一能做到这一点的量化方法，但其他方法无法像 PQ 那样有效地减少内存大小。我们实际上可以计算 PQ 和其他方法的内存使用量和量化操作复杂度，如下所示：\n",
        "\n",
        "**kmeans = kD PQ = mk^\\*D^\\* = k^{1/m}D**\n",
        "\n",
        "我们知道**D**代表输入向量的维数，但**k**和**m**可能是新的。**k**表示将用于表示向量的质心（或*代码）的**总数。*m表示我们将*向量分割成的子向量***的**数量（稍后会详细介绍）。\n",
        "\n",
        "*（“代码”是指我们向量的量化表示）*\n",
        "\n",
        "![使用 k=2048 和 m=8 时的内存使用量（和复杂性）与维度。](https://cdn.sanity.io/images/vr8gru94/production/358bd5e61d3324c4484e35c402b4aa252e1e9eb9-1920x1080.png)\n",
        "\n",
        "使用 k=2048 和 m=8 时的内存使用量（和复杂性）与维度。\n",
        "\n",
        "这里的问题是，为了获得好的结果，建议的**k**值是**2048**（211）或更大[1]。给定向量维度**D**，没有 PQ 的聚类会给我们带来*非常高的*内存需求和复杂性：\n",
        "\n",
        "|  手术  |                 内存和复杂性                 |\n",
        "| :----: | :------------------------------------------: |\n",
        "| k-均值 |            kD = 2048×128 = 262144            |\n",
        "| 质子Q  | mk*D* = (k^(1/m))×D = (2048^(1/8))×128 = 332 |\n",
        "\n",
        "给定 m 值为**8**，PQ 的等效内存使用和*分配*复杂性显着降低 - 由于将向量*分块为子向量*，并且将子量化过程应用于那些较小的维度**k\\***和**D\\***，等于**k/m**和分别为**D/m**。\n",
        "\n",
        "第二个重要因素是量化器训练。量化器需要比 k 大几倍的数据集才能进行有效训练，即*没有乘积子量化*。\n",
        "\n",
        "使用子量化器，我们只需要 k* 的几个倍数**（即 k/m）** ——这仍然是一个很大的数字——但可以显着减少。\n",
        "\n",
        "## 产品量化的工作原理\n",
        "\n",
        "让我们来研究一下 PQ 的逻辑。我们通常会有许多向量（长度都相等）——但为了简单起见，我们将在示例中使用单个向量。\n",
        "\n",
        "简而言之，PQ 的过程是：\n",
        "\n",
        "- 取一个大的高维向量，\n",
        "- 将其分割成大小相等的块——我们的子向量，\n",
        "- 将每个子向量分配给其最近的*质心*（也称为再现/重建值），\n",
        "- 用唯一的 ID 替换这些质心值 - 每个 ID 代表一个质心\n",
        "\n",
        "![](https://d33wubrfki0l68.cloudfront.net/af00b6764682bea50979e2285c0380f99e06466e/48940/images/product-quantization-6.mp4)\n",
        "\n",
        "PQ 的高级视图：采用一个大向量，将其分割为子向量，将每个子向量分配给其最近的质心值，并用其唯一的 ID 替换质心值 - 生成一个微小的 ID 向量。\n",
        "\n",
        "在该过程结束时，我们将需要*大量内存的高维向量减少为需要**很少内存*的微小 ID 向量。\n",
        "\n",
        "我们的向量的长度为 D 12。我们首先将该向量拆分为**m**个子向量，如下所示：\n",
        "\n",
        "![这里我们将高维向量 x 分成几个子向量 u_j。](https://cdn.sanity.io/images/vr8gru94/production/76b14448d700a0d7c3e52abce3cba6464a13237f-1920x1020.png)\n",
        "\n",
        "这里我们将高维向量 x 分成几个子向量 u_j。"
      ],
      "metadata": {
        "id": "EI7pSmUGMSGs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "x = [1, 8, 3, 9, 1, 2, 9, 4, 5, 4, 6, 2]"
      ],
      "metadata": {
        "id": "aXVVrOIQMT-I"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "m = 4\n",
        "D = len(x)\n",
        "# ensure D is divisable by m\n",
        "assert D % m == 0\n",
        "# length of each subvector will be D / m (D* in notation)\n",
        "D_ = int(D / m)"
      ],
      "metadata": {
        "id": "Ue6BA9xfN0rh"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# now create the subvectors\n",
        "u = [x[row:row+D_] for row in range(0, D, D_)]\n",
        "u"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DHSmsY8uN2b_",
        "outputId": "dd987a42-8d75-4bbe-ca8b-583db30bb845"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[[1, 8, 3], [9, 1, 2], [9, 4, 5], [4, 6, 2]]"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**我们可以通过每个子向量的位置j**来引用它。\n",
        "\n",
        "对于下一部分，请考虑聚类。作为一个随机示例，给定大量向量，我们可以说*“我想要三个聚类”*，然后我们优化这些聚类质心，根据每个向量最近的质心将向量分为*三个类别。*\n",
        "\n",
        "对于 PQ，我们做同样的事情，但有一点细微的差别。每个子向量空间（子空间）都分配有自己的一组聚类 - 因此我们生成的是一组跨多个子空间的聚类算法。"
      ],
      "metadata": {
        "id": "XswfZMFuN9ba"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "k = 2**5\n",
        "assert k % m == 0\n",
        "k_ = int(k/m)\n",
        "print(f\"{k=}, {k_=}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TdkpE0v-N4oA",
        "outputId": "61cdf386-1b1c-48dd-f9e1-09f62b143dd8"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "k=32, k_=8\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from random import randint\n",
        "\n",
        "c = []  # our overall list of reproduction values\n",
        "for j in range(m):\n",
        "    # each j represents a subvector (and therefore subquantizer) position\n",
        "    c_j = []\n",
        "    for i in range(k_):\n",
        "        # each i represents a cluster/reproduction value position *inside* each subspace j\n",
        "        c_ji = [randint(0, 9) for _ in range(D_)]\n",
        "        c_j.append(c_ji)  # add cluster centroid to subspace list\n",
        "    # add subspace list of centroids to overall list\n",
        "    c.append(c_j)"
      ],
      "metadata": {
        "id": "lJpLLY1iOAu_"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "我们的每个子向量将被分配给这些质心之一。在 PQ 术语中，这些质心称为*再现值，并由***cⱼ,ᵢ**表示，其中**j**是我们的子向量标识符，**i**标识所选的质心（*每个子向量空间 j**有* ***k\\**** 个质心） 。\n",
        "\n",
        "![我们的子向量被替换为特定的质心向量，然后可以用特定于该质心向量的唯一 ID 替换。](https://cdn.sanity.io/images/vr8gru94/production/05bc75a87d7a07b8d844d13a1b954458e71431d2-1920x1080.png)\n",
        "\n",
        "我们的子向量被替换为特定的质心向量，然后可以用特定于该质心向量的唯一 ID 替换。"
      ],
      "metadata": {
        "id": "1J3eRtVEOI-T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def euclidean(v, u):\n",
        "    distance = sum((x - y) ** 2 for x, y in zip(v, u)) ** .5\n",
        "    return distance\n",
        "\n",
        "def nearest(c_j, u_j):\n",
        "    distance = 9e9\n",
        "    for i in range(k_):\n",
        "        new_dist = euclidean(c_j[i], u_j)\n",
        "        if new_dist < distance:\n",
        "            nearest_idx = i\n",
        "            distance = new_dist\n",
        "    return nearest_idx"
      ],
      "metadata": {
        "id": "wpBfuil9OC5j"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ids = []\n",
        "for j in range(m):\n",
        "    i = nearest(c[j], u[j])\n",
        "    ids.append(i)\n",
        "ids"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XrIipy1QONEy",
        "outputId": "ff94f495-5125-4e0b-ae91-024449764f88"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[1, 7, 4, 3]"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "当我们用 PQ 处理一个向量时，它被分成我们的子向量，然后处理这些子向量并分配给它们最近的（子）簇质心（再现值）。\n",
        "\n",
        "我们没有存储由**D\\***维质心表示的量化向量，而是用质心 ID 替换它。每个质心**cⱼ,ᵢ**都有自己的 ID，稍后可以通过我们的密码本**c**将这些 ID 值映射回完整的质心。"
      ],
      "metadata": {
        "id": "kLB51RXuOnwE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "q = []\n",
        "for j in range(m):\n",
        "    c_ji = c[j][ids[j]]\n",
        "    q.extend(c_ji)"
      ],
      "metadata": {
        "id": "qwzyMfzvOOJs"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "q"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JzTcEy2cOrf6",
        "outputId": "3f69c1fe-7323-42f9-f0e3-b1578934044d"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[5, 9, 5, 7, 3, 2, 6, 5, 4, 1, 9, 0]"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "这样，我们就将 12 维向量压缩为 4 维 ID 向量。为了简单起见，我们在这里使用了小维度，因此这种技术的好处可能并不明显。\n",
        "\n",
        "让我们从原来的 8 位整数的 12 维向量切换到更实际的 32 位浮点数的 128 维向量（我们将在下一节中使用）。压缩为仅包含*八个*维度的 8 位整数向量后，我们可以在性能上找到良好的平衡。\n",
        "\n",
        "**原始：128×32 = 4096 量化：8×8 = 64**\n",
        "\n",
        "这是一个很大的区别——64x！\n",
        "\n",
        "\n",
        "\n",
        "------\n",
        "\n",
        "## Faiss 中的 PQ 实施\n",
        "\n",
        "到目前为止，我们已经了解了 Python 中简单、可读的 PQ 实现背后的逻辑。实际上，我们不会使用它，因为它没有优化，而且我们已经在其他地方有了出色的实现。相反，我们会使用像[Faiss](https://www.pinecone.io/learn/series/faiss/)这样的库，或者像[Pinecone](https://www.pinecone.io/)这样的生产就绪服务。\n",
        "\n",
        "我们将了解如何在 Faiss 中构建 PQ 索引，甚至将了解如何将 PQ 与倒排文件 (IVF) 步骤相结合以提高搜索速度。\n",
        "\n",
        "在开始之前，我们需要获取数据。我们将使用 Sift1M 数据集。可以使用以下脚本下载并打开它：\n"
      ],
      "metadata": {
        "id": "JF_nmg4iOz9D"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import shutil\n",
        "import urllib.request as request\n",
        "from contextlib import closing\n",
        "\n",
        "# first we download the Sift1M dataset\n",
        "with closing(request.urlopen('ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz')) as r:\n",
        "    with open('sift.tar.gz', 'wb') as f:\n",
        "        shutil.copyfileobj(r, f)"
      ],
      "metadata": {
        "id": "cIWm4NxOO1TK"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tarfile\n",
        "\n",
        "# the download leaves us with a tar.gz file, we unzip it\n",
        "tar = tarfile.open('sift.tar.gz', \"r:gz\")\n",
        "tar.extractall()"
      ],
      "metadata": {
        "id": "oIeSwMBsPOWr"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "# now define a function to read the fvecs file format of Sift1M dataset\n",
        "def read_fvecs(fp):\n",
        "    a = np.fromfile(fp, dtype='int32')\n",
        "    d = a[0]\n",
        "    return a.reshape(-1, d + 1)[:, 1:].copy().view('float32')"
      ],
      "metadata": {
        "id": "r8hh7s-QPTAE"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# data we will search through\n",
        "wb = read_fvecs('./sift/sift_base.fvecs')  # 1M samples\n",
        "# also get some query vectors to search with\n",
        "xq = read_fvecs('./sift/sift_query.fvecs')\n",
        "# take just one query (there are many in sift_learn.fvecs)\n",
        "xq = xq[0].reshape(1, xq.shape[1])"
      ],
      "metadata": {
        "id": "TmkDUOQyPVbv"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "xq.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rEhcYc-IPYbO",
        "outputId": "885d21e0-f0a0-464a-8bd5-93a979ad6515"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1, 128)"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "wb.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-t5C4moOPZhg",
        "outputId": "cff02699-9c3a-48f8-df75-0ee12bb2f7e9"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1000000, 128)"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "现在让我们进入第一个索引：IndexPQ。\n",
        "\n",
        "### 指数PQ\n",
        "\n",
        "我们的第一个索引是使用 IndexPQ 的纯 PQ 实现。为了初始化索引，我们需要定义三个参数。"
      ],
      "metadata": {
        "id": "Fm1kM9vEPfQ2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import faiss\n",
        "\n",
        "D = wb.shape[1]\n",
        "m = 8\n",
        "assert D % m == 0\n",
        "nbits = 8  # number of bits per subquantizer, k* = 2**nbits\n",
        "index = faiss.IndexPQ(D, m, nbits)"
      ],
      "metadata": {
        "id": "Lv_VFmq7PgPV"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "我们有向量维度**D**，即我们想要将完整向量分割成的子向量的数量（我们必须断言**D**可以被**m**整除）。\n",
        "\n",
        "最后，我们包含 nbits 参数。这定义了每个子量化器可以使用的位数，我们可以将其转换为分配给每个子空间的质心数，即**k_ = 2\\*\\*nbits**。nbits为 11 时**，**每个子空间有**2048 个**质心。\n",
        "\n",
        "因为我们使用的是 PQ，它使用聚类——我们必须使用 xb 数据集来训练索引。"
      ],
      "metadata": {
        "id": "ySe61k0YPnh7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "index.is_trained"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tho9KaOjPlrs",
        "outputId": "b941d09d-8f69-43a2-97d2-9629eab8a7d2"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "False"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.train(wb)  # PQ training can take some time when using large nbits"
      ],
      "metadata": {
        "id": "yEfz9lLEP_ru"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "这样，我们就可以继续将向量添加到索引和搜索中。\n",
        "\n"
      ],
      "metadata": {
        "id": "-QPPcyn1QERF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dist, I = index.search(xq, k)\n",
        "print(I,dist)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hfa_itznQGHM",
        "outputId": "76397f93-9ebd-4b9f-9a37-cb2a5e60dab5"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[-1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1\n",
            "  -1 -1 -1 -1 -1 -1 -1 -1]] [[3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38\n",
            "  3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38\n",
            "  3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38\n",
            "  3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38\n",
            "  3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38\n",
            "  3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38 3.4028235e+38\n",
            "  3.4028235e+38 3.4028235e+38]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "index.search(xq, k)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JeuqCT00QKEM",
        "outputId": "2a74d4a9-28e9-4b6f-ba74-54b356f798b7"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "92.1 µs ± 62.6 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "根据我们的搜索，我们将返回前k个最接近的匹配项（与之前表示法中使用的k不同）。返回dist中的距离和I中的索引。\n",
        "\n",
        "我们可以将IndexPQ的召回性能与平面索引的召回性能进行比较，后者具有“完美”的召回率（得益于不压缩向量并执行详尽的搜索）。\n",
        "\n"
      ],
      "metadata": {
        "id": "zc6BvPDAQglZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "l2_index = faiss.IndexFlatL2(D)\n",
        "l2_index.add(wb)"
      ],
      "metadata": {
        "id": "MxnDcNosQhoW"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "l2_dist, l2_I = l2_index.search(xq, k)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "J_wmAdYnQmLD",
        "outputId": "0c288405-b0b8-4493-8a64-82cead4745a2"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 60.7 ms, sys: 848 µs, total: 61.6 ms\n",
            "Wall time: 65.1 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sum([1 for i in I[0] if i in l2_I])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ReGqzXr3QqqP",
        "outputId": "2c7cf197-0729-40a3-dfad-0fbd4463f556"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "如果我们愿意为了减少 PQ 的内存使用而牺牲完美的结果，我们会得到 50%，这是一个合理的召回率。平面搜索时间也减少到了 18%——我们可以在以后使用 IVF进一步改进这一点。\n",
        "\n",
        "较低的召回率是 PQ 的一个主要缺点。这可以通过使用较大的nbits值来抵消，但代价是搜索时间较慢且索引构建时间非常慢。然而，PQ 和 IVFPQ 索引都无法达到很高的召回率。如果需要更高的召回率，则应考虑另一个索引。\n",
        "\n",
        "在内存使用方面，IndexPQ 与我们的扁平索引相比如何？\n",
        "\n"
      ],
      "metadata": {
        "id": "RzRbGC_yQvs4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "def get_memory(index):\n",
        "    # write index to file\n",
        "    faiss.write_index(index, './temp.index')\n",
        "    # get file size\n",
        "    file_size = os.path.getsize('./temp.index')\n",
        "    # delete saved index\n",
        "    os.remove('./temp.index')\n",
        "    return file_size"
      ],
      "metadata": {
        "id": "2gGmzgxkQxUe"
      },
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# first we check our Flat index, this should be larger\n",
        "get_memory(l2_index)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YAQ2du-TQywp",
        "outputId": "afd02643-58cf-40ba-a4c8-aea6d527a505"
      },
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "512000045"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# now our PQ index\n",
        "get_memory(index)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R-4L5Pk3Q0HN",
        "outputId": "909b4114-123d-4b40-a252-a11bb2bc1baf"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "131158"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "简而言之，使用 IndexPQ 的内存使用量非常出色，内存减少了 98.4%。通过使用 IVF+PQ 索引，也可以将其中一些荒谬的性能优势转化为搜索速度。\n",
        "\n",
        "### 指数IVFPQ\n",
        "\n",
        "为了加快搜索时间，我们可以添加另一个步骤，即使用 IVF 索引，这将作为减少搜索中向量范围的初始粗略步骤。\n",
        "\n",
        "之后，我们像之前一样继续 PQ 搜索，但向量数量显着减少。由于最小化了我们的搜索范围，我们应该发现我们的搜索速度得到了极大的提高。\n",
        "\n",
        "让我们看看它是如何工作的。首先，我们像这样初始化 IVF+PQ 索引："
      ],
      "metadata": {
        "id": "OnMHX-GXQ8bZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "vecs = faiss.IndexFlatL2(D)\n",
        "\n",
        "nlist = 2048  # how many Voronoi cells (must be >= k* which is 2**nbits)\n",
        "nbits = 8  # when using IVF+PQ, higher nbits values are not supported\n",
        "index = faiss.IndexIVFPQ(vecs, D, nlist, m, nbits)\n",
        "print(f\"{2**nbits=}\")  # our value for nlist"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fnfUjJKzQ9Uz",
        "outputId": "c1be64a3-c3f4-436d-904d-dde4202410d1"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2**nbits=256\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "我们在这里有一个新参数，**nlist**定义了我们使用多少个 Voronoi 单元来对*已经*量化的 PQ 向量进行聚类（[在此处了解有关 IndexIVF 的更多信息](https://www.pinecone.io/learn/series/faiss/vector-indexes/)）。您可能会问，Voronoi单元到底是什么？这到底意味着什么？让我们可视化一些 2D“PQ 向量”：\n",
        "\n",
        "![显示我们重建的“PQ”向量的二维图表。 然而，实际上，我们永远不会对 2D 向量使用 PQ，因为没有足够的维数来分割子向量和子量化。](https://cdn.sanity.io/images/vr8gru94/production/9e0dbf3ff9d987e13fa99601f4fd0dbb88d78e68-1920x1080.png)\n",
        "\n",
        "显示我们重建的“PQ”向量的二维图表。然而，实际上，我们永远不会对 2D 向量使用 PQ，因为没有足够的维数来分割子向量和子量化。\n",
        "\n",
        "让我们添加一些 Voronoi 单元：\n",
        "\n",
        "![二维图表显示了我们的量化“PQ”向量，这些向量现已通过 IVF 分配给不同的 Voronoi 细胞。](https://cdn.sanity.io/images/vr8gru94/production/cac8f97499227e918bca966c57b8c00591a81f3f-1920x1080.png)\n",
        "\n",
        "二维图表显示了我们的量化“PQ”向量，这些向量现已通过 IVF 分配给不同的 Voronoi 细胞。\n",
        "\n",
        "在较高的层次上，它们只是一组分区。相似的向量被分配给不同的分区（或*单元格*），当涉及到搜索时，我们引入查询向量 xq 并将我们的搜索限制到最近的单元格：\n",
        "\n",
        "![IVF 允许我们将搜索限制为仅分配到附近细胞的向量。 洋红色点是我们的查询向量 xq。 现在让我们继续进行训练和搜索，看看我们的搜索速度和召回率如何。](https://cdn.sanity.io/images/vr8gru94/production/5e002882007c9546c5a96cb79ac300bf3226ff7a-1920x1080.png)\n",
        "\n",
        "IVF 允许我们将搜索限制为仅分配到附近细胞的向量。洋红色点是我们的查询向量 xq。现在让我们继续进行训练和搜索，看看我们的搜索速度和召回率如何。"
      ],
      "metadata": {
        "id": "olwpBm7JRKcH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "index.train(wb)"
      ],
      "metadata": {
        "id": "yKVf72bxRIu9"
      },
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index.add(wb)"
      ],
      "metadata": {
        "id": "DsCApRwMRddp"
      },
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dist, I = index.search(xq, k)\n",
        "print(I,dist)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "L6rD_hztRd3V",
        "outputId": "aaddbf6a-96fe-48b1-e3cc-11f13858b30f"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[562594 562167 170166 225116 746238 930567 455537  36538 580841  49874\n",
            "  312682 170996 562343 732041 742261 541845 989762  16429 392032 957845\n",
            "  562556 403442 629490 670103 779712 377461 839679 619660 565419 465294\n",
            "  656997 490859]] [[60581.395 61242.836 65607.59  66642.48  67303.734 67706.42  68118.09\n",
            "  68135.836 69565.13  69808.01  69860.92  69975.11  70398.85  70417.31\n",
            "  70422.484 70432.59  70508.125 70519.76  71256.53  71877.445 71969.77\n",
            "  72802.58  73408.125 73466.72  73950.33  74384.95  74463.836 75076.29\n",
            "  75134.94  75307.73  75381.67  75428.6  ]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "index.search(xq, k)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EqfsjmHbRfHE",
        "outputId": "ef9a99f4-cc79-4dae-c12a-eefc57385763"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "94.2 µs ± 1.67 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "搜索时间快如闪电，为 86.3μs，但召回率与我们的 IndexPQ 相比显着下降（50% 降至 34%）。给定相同的参数，IndexPQ 和 IndexIVFPQ*应该*能够获得相同的召回性能。\n",
        "\n",
        "在这种情况下，提高召回率的秘诀是提高**nprobe**参数——它告诉我们*有多少个*最近的 Voronoi 单元包含在我们的搜索范围中。\n",
        "\n",
        "![](https://d33wubrfki0l68.cloudfront.net/55d7c99f61b9d00f795f15d30340af276758f139/9d76b/images/product-quantization-11-ivf-nprobe.mp4)\n",
        "\n",
        "*二维图表显示了我们的量化“PQ”向量，这些向量现已通过 IVF 分配给不同的 Voronoi 细胞。*\n",
        "\n",
        "在一种极端情况下，我们可以通过将 nprobe 设置为**nlist**值来包含所有单元格 - 这将返回最大可能的召回率。\n",
        "\n",
        "当然，我们不想包含所有单元格。这将使我们的 IVF 索引毫无意义，因为这样我们就不会限制我们的搜索范围（使其相当于平面索引）。相反，我们找到了实现这种召回性能的最低**nprobe值。**"
      ],
      "metadata": {
        "id": "FkKn4wpxRhxj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "index.nprobe = 2048  # this is searching all Voronoi cells\n",
        "dist, I = index.search(xq, k)\n",
        "sum([1 for i in I[0] if i in l2_I])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qCMuedywSkt1",
        "outputId": "a6593e82-a50e-4293-ea9b-9969902f74f6"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "9"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.nprobe = 2\n",
        "dist, I = index.search(xq, k)\n",
        "sum([1 for i in I[0] if i in l2_I])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QXMYPUV-SmB5",
        "outputId": "4e4dd3b7-4bdf-4c9d-eef9-aaaeb8270b2c"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "6"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.nprobe = 48\n",
        "dist, I = index.search(xq, k)\n",
        "sum([1 for i in I[0] if i in l2_I])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Onx4JB5TS0os",
        "outputId": "30fc9fb8-8de4-4cb8-8971-bb7a012c50f6"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "9"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**当 nprobe 为 48 时，我们实现了 52% 的最佳召回率（如nprobe == 2048**所示），同时最小化我们的搜索范围（从而最大化搜索速度）。\n",
        "\n",
        "通过添加 IVF 步骤，我们将搜索时间从 IndexPQ 的 1.49 毫秒大幅减少到 IndexIVFPQ 的 0.09 毫秒。得益于我们的 PQ 向量，我们将其与比 Flat 索引低 96% 的极小内存使用量配对。\n",
        "\n",
        "总而言之，IndexIVFPQ*极大地减少*了内存使用量（尽管比 IndexPQ 略大，分别为 9.2MB 和 6.5MB）以及闪电般的搜索速度，同时保持了 50% 左右的合理召回率。\n",
        "\n",
        "------\n",
        "\n",
        "这就是本文的内容！我们已经介绍了乘积量化 (PQ) 背后的直觉，以及它如何压缩我们的索引并实现令人难以置信的高效内存使用。\n",
        "\n",
        "我们将 Faiss IndexPQ 实现放在一起，并测试了搜索时间、召回率和内存使用情况，然后通过使用 IndexIVFPQ 将其与 IVF 索引配对来进一步优化索引。\n",
        "\n",
        "|              | 平L2 | 质子Q | IVFPQ |\n",
        "| :----------: | :--: | :---: | :---: |\n",
        "| 记起 （％）  | 100  |  50   |  52   |\n",
        "| 速度（毫秒） | 8.26 | 1.49  | 0.09  |\n",
        "|  内存（MB）  | 256  |  6.5  |  9.2  |\n",
        "\n",
        "我们的测试结果显示出令人印象深刻的内存压缩和搜索速度，以及合理的召回分数。\n",
        "\n",
        "如果您有兴趣了解有关搜索中可用的各种索引的更多信息，包括有关 IVF 索引的更多详细信息，请阅读我们[关于相似性搜索的最佳索引的文章](https://www.pinecone.io/learn/series/faiss/vector-indexes/)。\n",
        "\n",
        "## 参考\n",
        "\n",
        "- [1] H Jégou 等人，[最近邻搜索的乘积量化](https://inria.hal.science/inria-00514462v2/document)(2010)\n",
        "\n",
        "- [Jupyter 笔记本](https://github.com/pinecone-io/examples/tree/master/learn/search/faiss-ebook/product-quantization)\n",
        "\n",
        "  "
      ],
      "metadata": {
        "id": "qj5Cfw4-TVje"
      }
    }
  ]
}