GPU K-means示例

#include <cstdio>
#include <vector>

#include <faiss/Clustering.h>
#include <faiss/utils.h>
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/StandardGpuResources.h>

// 在超立方体中生成一些随机向量（在CPU上操作）
std::vector<float> makeRandomVecs(size_t numVecs, int dim) {
  std::vector<float> vecs(numVecs * dim);
  faiss::float_rand(vecs.data(), vecs.size(), 1);
  return vecs;
}

int main(int argc, char** argv) {
  // 默认会预留18%的GPU内存用于临时计算工作；
  // 可以根据实际需要调整这个比例，
  // 也可以自定义GpuResources对内存进行不同方式的管理。
  faiss::gpu::StandardGpuResources res;

  int dim = 128; // 向量维度
  int numberOfEMIterations = 20; // EM迭代次数
  size_t numberOfClusters = 20000; // 聚类中心个数
  size_t numVecsToCluster = 5000000; // 参与聚类的数据向量总数

  // 生成一批随机向量；注意：这里生成的数据是在CPU上
  std::vector<float> vecs = makeRandomVecs(numVecsToCluster, dim);

  faiss::gpu::GpuIndexFlatConfig config;
  config.device = 0;            // 使用第0号GPU（默认）
  config.useFloat16 = false;    // 不使用float16精度（默认）
  faiss::gpu::GpuIndexFlatL2 index(&res, dim, config);

  faiss::ClusteringParameters cp;
  cp.niter = numberOfEMIterations;
  cp.verbose = true; // 输出每轮迭代的统计信息

  // 如果需要球面（spherical）k-means，请使用GpuIndexFlatIP，并设置cp.spherical = true

  // :::note
  // Faiss默认情况下，每个聚类中心采样256个向量，以避免在聚类中心数较少、向量数巨大的情况下资源耗尽。
  // 例如，numberOfClusters = 1000，而numVecsToCluster = 1000000时，
  // 实际只会采样256000个向量参与聚类。
  // 你可以通过设置cp.max_points_per_centroid
  // 覆盖此采样逻辑，使每个聚类中心使用自定义数量的数据：
  // cp.max_points_per_centroid = ((numVecsToCluster + numberOfClusters - 1) / numberOfClusters);
  // :::

  faiss::Clustering kMeans(dim, numberOfClusters, cp);

  // 开始执行聚类
  kMeans.train(numVecsToCluster, vecs.data(), index);

  // kMeans.centroids 包含最终得到的聚类中心（数据保存在CPU上）
  printf("centroid 3 dim 6 is %f\n", kMeans.centroids[3 * dim + 6]);

  return 0;
}