faiss库中ivf-sq（ScalarQuantizer，标量量化）代码解读-2

文件ScalarQuantizer.h

主要介绍这里面的枚举以及一些函数内容：QuantizerType、RangeStat、ScalarQuantizer、train、compute_codes、decode、SQuantizer、FlatCodesDistanceComputer、get_distance_computer、select_InvertedListScanner

QuantizerType

量化类型中包含一些枚举类型，代码内容如下：

enum QuantizerType {QT_8bit,         ///< 8 bits per componentQT_4bit,         ///< 4 bits per componentQT_8bit_uniform, ///< same, shared range for all dimensionsQT_4bit_uniform,QT_fp16,QT_8bit_direct, ///< fast indexing of uint8sQT_6bit,        ///< 6 bits per componentQT_bf16,QT_8bit_direct_signed, ///< fast indexing of signed int8s ranging from///< [-128 to 127]};

QuantizerType 枚举内容解析

枚举值	描述	适用场景
QT_8bit	每个分量用 8-bit 表示，分量间有独立范围（非均匀量化）。	精度较高，适用于一般场景。
QT_4bit	每个分量用 4-bit 表示，分量间有独立范围（非均匀量化）。	精度较低但压缩率高，适用于存储敏感的大规模数据。
QT_8bit_uniform	每个分量用 8-bit 表示，但所有分量共享一个范围（均匀量化）。	适合数据分布较均匀的场景，计算简单，存储开销低。
QT_4bit_uniform	每个分量用 4-bit 表示，所有分量共享一个范围（均匀量化）。	存储极为节省，但精度较低，适合快速粗筛。
QT_fp16	每个分量用 16-bit 浮点数（half-precision floating-point） 表示。	高精度场景，兼顾压缩率和精度，但存储开销大于 8-bit 和 4-bit。
QT_8bit_direct	直接使用无符号 uint8 值，无需量化，直接存储整数。	数据本身就是无符号 8-bit 整数，例如像素值。
QT_6bit	每个分量用 6-bit 表示，比 8-bit 压缩但比 4-bit 精度更高。	存储与精度需求的折中选择，适合数据分布复杂但存储受限的场景。
QT_bf16	使用 bfloat16（16-bit 浮点数） 表示，范围与精度介于 fp16 和 uint8 之间。	适用于高动态范围但对精度要求不高的场景，例如部分深度学习推理中。
QT_8bit_direct_signed	直接使用有符号 int8 值，无需量化，范围为 [-128, 127]。	数据本身就是有符号的整数，例如分量可能包含负值的数据。

为什么设置为这些类型？

支持不同的精度需求

• 数据量化需要在存储成本与精度之间找到平衡，不同类型的量化方式提供了不同的压缩率和精度。
  • 8-bit 和 4-bit：提供了常见的两种量化方式，适合大多数存储需求。
  • 6-bit：为 8-bit 和 4-bit 之间提供折中选择。
  • fp16 和 bf16：用于精度较高的场景，支持浮点值表示。
    适配多种数据分布
  • 有些数据在不同维度的范围不同（非均匀分布），需要 非均匀量化（如 QT_8bit 和 QT_4bit）。
  • 对于范围均匀的数据，可以使用 均匀量化（如 QT_8bit_uniform 和 QT_4bit_uniform），进一步减少存储复杂性。
    提高存储效率
  • QT_4bit 和 QT_4bit_uniform：只用 4-bit 表示每个分量，适合存储空间受限的大规模向量。
  • QT_6bit：比 8-bit 进一步压缩，但比 4-bit 提供更高精度。
    兼容特殊数据类型
  • QT_8bit_direct 和 QT_8bit_direct_signed：直接存储原始值，不需要量化过程，适合已经量化好的数据或离散整数值。

RangeStat

代码内容如下：

enum RangeStat {RS_minmax,    ///< [min - rs*(max-min), max + rs*(max-min)]RS_meanstd,   ///< [mean - std * rs, mean + std * rs]RS_quantiles, ///< [Q(rs), Q(1-rs)]RS_optim,     ///< alternate optimization of reconstruction error
};

枚举值	描述	用途与适用场景
RS_minmax	根据向量的最小值和最大值生成范围： [ min - rs * (max - min), max + rs * (max - min) ]。	- 数据分布均匀的场景 - 快速计算，但对异常值较为敏感。
RS_meanstd	根据向量的均值和标准差生成范围： [ mean - std * rs, mean + std * rs ]。	- 考虑数据集中趋势和离散程度 - 对异常值有一定鲁棒性，适合常规分布数据。
RS_quantiles	根据分位数生成范围： [ Q(rs), Q(1-rs) ]。	- 非参数方法，适合分布不均的数据 - 有效减少异常值的影响（如 5%-95% 分位数）。
RS_optim	基于优化误差动态调整范围。	- 重建精度要求高的场景 - 适合对误差敏感的任务，可能效率较低，但结果更精确。

• rs 是范围缩放因子（Range Scale）。rs是用户定义的参数，根据具体场景和需求选择合格的值。
• mean：向量分量的平均值（Mean）。
• std：向量分量的标准差。

参数	含义	作用
rs	范围缩放因子，控制范围的扩展程度。	调节范围宽度。rs 越大，范围越宽，适合分布离散的情况；rs 越小，范围越窄，适合分布集中的情况。
mean	均值，表示数据的集中趋势。	用作范围的中心值，减少偏移误差。
std	标准差，表示数据的离散程度。	用于计算范围宽度，离散程度越大，生成的范围越宽。

ScalarQuantizer

用于将浮点转化为相对应的量化结果
头文件C语言内容：

#ifndef SCALAR_QUANTIZER_H
#define SCALAR_QUANTIZER_H#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>/* Quantizer 类型枚举 */
typedef enum {QT_8bit,            /* 8 位每分量 */QT_4bit,            /* 4 位每分量 */QT_8bit_uniform,    /* 每分量 8 位，范围统一 */QT_4bit_uniform,    /* 每分量 4 位，范围统一 */QT_fp16,            /* 16 位浮点 */QT_bf16,            /* 16 位 bfloat */QT_8bit_direct,     /* 无需额外计算的直接 8 位编码 */QT_8bit_direct_signed, /* 带符号的直接 8 位编码 */QT_6bit             /* 每分量 6 位 */
} QuantizerType;/* ScalarQuantizer 数据结构 */
typedef struct {size_t d;            /* 输入向量的维度 */size_t code_size;    /* 每个向量的编码大小（字节数） */int bits;            /* 每分量的位数 */QuantizerType qtype; /* 量化器类型 */
} ScalarQuantizer;/* 初始化函数 */
void ScalarQuantizer_init(ScalarQuantizer* sq, size_t d, QuantizerType qtype);/* 设置 derived sizes */
void ScalarQuantizer_set_derived_sizes(ScalarQuantizer* sq);/* 打印量化器信息 */
void ScalarQuantizer_print_info(const ScalarQuantizer* sq);#endif /* SCALAR_QUANTIZER_H */

实现部分：

#include "ScalarQuantizer.h"/* 初始化 ScalarQuantizer */
void ScalarQuantizer_init(ScalarQuantizer* sq, size_t d, QuantizerType qtype) {if (sq == NULL) {fprintf(stderr, "Error: ScalarQuantizer pointer is NULL.\n");exit(EXIT_FAILURE);}sq->d = d;sq->qtype = qtype;sq->code_size = 0; /* 初始化为 0 */sq->bits = 0;      /* 初始化为 0 */ScalarQuantizer_set_derived_sizes(sq);
}/* 设置 code_size 和 bits */
void ScalarQuantizer_set_derived_sizes(ScalarQuantizer* sq) {if (sq == NULL) {fprintf(stderr, "Error: ScalarQuantizer pointer is NULL.\n");exit(EXIT_FAILURE);}switch (sq->qtype) {case QT_8bit:case QT_8bit_uniform:case QT_8bit_direct:case QT_8bit_direct_signed:sq->code_size = sq->d;sq->bits = 8;break;case QT_4bit:case QT_4bit_uniform:sq->code_size = (sq->d + 1) / 2; /* 每两个分量占 1 字节 */sq->bits = 4;break;case QT_6bit:sq->code_size = (sq->d * 6 + 7) / 8; /* 位对齐处理 */sq->bits = 6;break;case QT_fp16:case QT_bf16:sq->code_size = sq->d * 2; /* 每分量 2 字节 */sq->bits = 16;break;default:fprintf(stderr, "Error: Unsupported QuantizerType.\n");exit(EXIT_FAILURE);}/* 检查初始化结果 */if (sq->code_size == 0 || sq->bits == 0) {fprintf(stderr, "Error: Invalid derived sizes.\n");exit(EXIT_FAILURE);}
}/* 打印 ScalarQuantizer 信息 */
void ScalarQuantizer_print_info(const ScalarQuantizer* sq) {if (sq == NULL) {fprintf(stderr, "Error: ScalarQuantizer pointer is NULL.\n");return;}printf("ScalarQuantizer Info:\n");printf("  Dimension (d): %zu\n", sq->d);printf("  Code Size (bytes per vector): %zu\n", sq->code_size);printf("  Bits per component: %d\n", sq->bits);printf("  Quantizer Type: %d\n", sq->qtype);
}

关键点说明

1. ScalarQuantizer 数据结构：

• 包含 d（维度）、code_size（编码大小）、bits（每分量位数）和 qtype（量化器类型）。

2. 初始化函数：

• ScalarQuantizer_init 接收维度和量化器类型作为输入，并调用 ScalarQuantizer_set_derived_sizes 设置计算派生值。

3. 错误处理：

• 添加了对空指针的检查以及对非法 QuantizerType 的错误处理。

4. 位对齐逻辑：

• 在 QT_6bit 和 QT_4bit 的计算中，处理了位对齐问题。

train、compute_codes和decode

上述函数的作用参考上一篇链接：标量量化
这里我定义的ScalarQuantizer结构体内容如下：

typedef struct ScalarQuantizer{size_t d;              /* 输入向量的维度 */size_t code_size;      /* 每个向量的编码大小（字节数） */int bits;              /* 每分量的位数 */QuantizerType qtype;   /* 量化器类型 *//* 训练参数 */float rangestat;       /* 范围统计参数 */float rangestat_arg;   /* 范围统计扩展参数 */int n_centroids;       /* 聚类中心数量 *//* 训练结果 */float* trained;        /* 存储训练的量化中心 */float* centroids;      /* 非均匀量化时的聚类中心 *//* 并行和状态跟踪 */int num_threads;       /* 并行计算的线程数量 */size_t progress;       /* 当前训练/编码的进度 *//* 高级量化支持 */int* per_dim_bits;     /* 每维分量的位数，支持动态调整 */float* residuals;      /* 残差向量存储 */int residual_levels;   /* 残差量化层数 */
} ScalarQuantizer;

然后将训练train分为3个部分：ScalarQuantizer_train（QT类4、6、8位）、train_Uniform（均匀量化训练 ）和train_NonUniform（非均匀量化训练），其中train_Uniform和train_NonUniform在代码中呈现。

void ScalarQuantizer::train(size_t n, const float* x) {int bit_per_dim = qtype == QT_4bit_uniform ? 4: qtype == QT_4bit                 ? 4: qtype == QT_6bit                 ? 6: qtype == QT_8bit_uniform         ? 8: qtype == QT_8bit                 ? 8: -1;switch (qtype) {case QT_4bit_uniform:case QT_8bit_uniform:train_Uniform(rangestat,rangestat_arg,n * d,1 << bit_per_dim,x,trained);break;case QT_4bit:case QT_8bit:case QT_6bit:train_NonUniform(rangestat,rangestat_arg,n,d,1 << bit_per_dim,x,trained);break;case QT_fp16:case QT_8bit_direct:case QT_bf16:case QT_8bit_direct_signed:// no training necessarybreak;}
}