TensorFlow Recommender含ScaNN索引大模型保存内存瓶颈求解

针对你用TFRS构建检索模型、ScaNN实现索引，在500万+数据集下保存模型时内存爆仓的问题，给你几个实用的解决方案：

1. 延迟ScaNN索引构建，仅保存模型权重 #

默认情况下，ScaNN层训练完成后会预构建完整索引并随模型一起保存，这会导致大索引直接占满内存。你可以在保存前重置索引，只保存模型结构和权重，等到推理阶段再重新构建索引：

# 保存前重置ScaNN索引
scann_layer.reset_index()

# 保存模型
tf.saved_model.save(model, "./saved_model")

# 加载模型后，重新用训练好的embedding构建索引
loaded_model = tf.saved_model.load("./saved_model")
# 假设你有包含embedding和id的数据集dataset
loaded_model.scann_layer.index_from_dataset(
    dataset.map(lambda x: (x['user_embedding'], x['user_id']))
)

2. 调整ScaNN索引参数降低内存占用 #

通过减少聚类数、缩小聚类训练样本量，可以大幅降低ScaNN索引的内存开销，同时对检索精度的影响可控：

# 构建ScaNN层时调整参数
scann_layer = tfrs.layers.factorized_top_k.ScaNN(
    query_model=query_model,
    k=10,
    num_leaves=2000,  # 减少聚类数，默认可能是数千到上万
    training_sample_size=100000,  # 用10万样本训练聚类中心，而非全量500万
    distance_measure="dot_product"
)

注意：num_leaves和training_sample_size需要根据你的数据集规模调整，聚类数过小可能影响检索速度，样本量过小可能降低聚类质量。

3. 分批次构建并保存ScaNN索引 #

如果必须保存预构建的索引，可以分批次处理embedding，避免一次性加载全量数据到内存：

# 分批次生成索引
batch_size = 100000
dataset_batches = dataset.batch(batch_size)

# 初始化索引
scann_layer.index_dataset(
    dataset_batches.take(1).map(lambda x: (x['embedding'], x['id']))
)

# 逐批次添加剩余数据
for batch in dataset_batches.skip(1):
    embeddings, ids = batch['embedding'], batch['id']
    scann_layer.add(embeddings, ids)

# 再保存模型
tf.saved_model.save(model, "./saved_model_with_index")

4. 升级TF/TFRS版本或更换保存方式 #

TF 2.9.1的TFRS可能存在ScaNN保存时的内存泄漏问题，升级到TF 2.12+及对应的TFRS版本（比如TFRS 0.7+），官方修复了不少内存相关的bug。另外，尝试用tf.keras.models.save_model()替代tf.saved_model.save()，部分场景下前者的内存管理更高效：

tf.keras.models.save_model(model, "./keras_saved_model")

5. 临时扩容内存或启用Swap #

如果云端虚拟机允许，临时把内存扩容到32GB以上可以直接解决问题；如果无法扩容，可启用Swap分区作为内存补充（速度会变慢，但能避免进程被终止）：

# 在容器宿主机执行以下命令创建16GB Swap
fallocate -l 16G /swapfile
chmod 600 /swapfile
mkswap /swapfile
swapon /swapfile

内容的提问来源于stack exchange，提问作者Pysnek313