Torch-RecHub模型部署教程：ONNX导出与向量索引，从训练到生产环境一键部署

【免费下载链接】torch-rechub A Lighting Pytorch Framework for Recommendation Models, Easy-to-use and Easy-to-extend. 项目地址: https://gitcode.com/gh_mirrors/to/torch-rechub

Torch-RecHub是一个基于PyTorch的推荐系统框架，提供了从模型训练到部署的完整解决方案。本教程将详细介绍如何使用Torch-RecHub实现模型的ONNX导出与量化，以及向量索引的构建与查询，帮助你轻松实现推荐系统从训练到生产环境的一键部署。

推荐系统部署全流程概览

在推荐系统中，模型部署是连接科研与业务的关键桥梁。一个完整的推荐系统部署流程通常包括模型导出、优化、向量索引构建和在线服务等环节。Torch-RecHub提供了一套完整的工具链，简化了这一过程。

上图展示了Torch-RecHub推荐系统部署的完整流程，从模型训练到最终的在线服务，涵盖了ONNX导出、量化优化和向量索引等关键步骤。

环境准备与安装

在开始部署之前，我们需要先安装Torch-RecHub及其相关依赖。首先，通过以下命令克隆项目仓库：

git clone https://gitcode.com/gh_mirrors/to/torch-rechub
cd torch-rechub

然后，安装ONNX导出和向量索引所需的依赖：

# 安装ONNX相关依赖
pip install "torch-rechub[onnx]"

# 安装向量索引相关依赖
pip install "torch-rechub[annoy,faiss,milvus]"

注意：如果需要在GPU环境下运行ONNX推理，请安装onnxruntime-gpu替代默认的onnxruntime。对于Milvus，还需要额外启动Milvus服务，请参考官方文档进行安装配置。

ONNX导出：跨平台推理的关键一步

ONNX（Open Neural Network Exchange）是一种开放的神经网络模型格式，支持多种深度学习框架和推理引擎。将模型导出为ONNX格式，可以实现跨平台、跨框架的推理部署。

ONNX导出与量化流程

Torch-RecHub提供了便捷的ONNX导出和量化工具，支持从训练器直接导出模型，并进行INT8/FP16量化优化。

不同类型模型的导出方法

Torch-RecHub支持多种类型的推荐模型导出，包括CTR排序模型、Matching召回模型和多任务模型等。

1. CTR排序模型导出

对于DeepFM、DIN等排序模型，可以直接使用CTRTrainer的export_onnx方法：

from torch_rechub.trainers import CTRTrainer

# 假设已经训练好了一个CTR模型
# trainer.fit(train_dl, val_dl)

# 导出ONNX模型
trainer.export_onnx("deepfm.onnx")

2. Matching召回模型导出

对于DSSM、YoutubeDNN等双塔召回模型，通常需要分别导出用户塔和物品塔：

from torch_rechub.trainers import MatchTrainer

# 导出用户塔
trainer.export_onnx("user_tower.onnx", mode="user")

# 导出物品塔
trainer.export_onnx("item_tower.onnx", mode="item")

3. 多任务模型导出

对于MMOE、PLE等多任务模型，可以使用MTLTrainer的export_onnx方法：

from torch_rechub.trainers import MTLTrainer

# 导出多任务模型
trainer.export_onnx("mmoe.onnx")

高级导出选项

Torch-RecHub还支持动态batch size、导出器选择等高级功能：

# 动态batch size导出
trainer.export_onnx(
    "model.onnx",
    dynamic_batch=True,  # 启用动态batch size
    onnx_export_kwargs={
        "dynamo": False,  # 使用legacy导出器
    },
)

ONNX量化：提升推理性能的关键优化

为了在生产环境中获得更好的性能，通常需要对ONNX模型进行量化优化。Torch-RecHub支持INT8动态量化和FP16转换两种方式。

INT8动态量化（推荐CPU环境）

INT8量化可以显著减小模型大小，提高CPU推理速度：

from torch_rechub.utils.quantization import quantize_model

# INT8动态量化
quantize_model(
    input_path="model_fp32.onnx",
    output_path="model_int8.onnx",
    mode="int8",
    per_channel=True,  # 启用per-channel量化
    reduce_range=True  # 缩小量化范围，提高稳定性
)

FP16转换（推荐GPU环境）

对于支持Tensor Core的GPU，FP16转换可以提高推理速度并减少显存占用：

# FP16转换
quantize_model(
    input_path="model_fp32.onnx",
    output_path="model_fp16.onnx",
    mode="fp16",
    keep_io_types=True  # 保持输入输出为FP32，提高兼容性
)

量化性能对比

Torch-RecHub提供了量化性能对比脚本，可以方便地比较不同量化方式的效果：

# 对比FP32和INT8性能
python examples/serving/benchmark_onnx_quantization.py --fp32 model_fp32.onnx --int8 model_int8.onnx

# 对比FP32和FP16性能（GPU）
python examples/serving/benchmark_onnx_quantization.py --fp32 model_fp32.onnx --fp16 model_fp16.onnx --provider CUDAExecutionProvider

向量索引：高效推荐召回的核心组件

在推荐系统中，向量索引是实现高效召回的关键。Torch-RecHub提供了统一的向量检索接口，支持Annoy、FAISS和Milvus三种主流的近似最近邻搜索库。

向量索引架构

Torch-RecHub采用Builder-Indexer模式设计向量索引组件，提供了一致的API接口，方便用户在不同检索后端之间切换。

三种索引后端的使用方法

1. Annoy索引（轻量级，内存友好）

Annoy是Spotify开源的近似最近邻搜索库，适合小规模数据和内存受限的场景：

from torch_rechub.serving import builder_factory

# 创建Annoy索引构建器
builder = builder_factory(
    "annoy",
    d=64,  # 向量维度
    metric="angular",  # 距离度量（余弦相似度）
    n_trees=50  # 树的数量，影响精度和速度
)

# 从嵌入向量构建索引并查询
with builder.from_embeddings(item_embeddings) as indexer:
    ids, distances = indexer.query(user_embeddings, top_k=10)
    indexer.save("annoy_index.ann")  # 保存索引

2. FAISS索引（高性能，适合中大规模数据）

FAISS是Meta开源的高性能相似性搜索库，支持多种索引类型，适合中大规模数据：

# 创建FAISS HNSW索引构建器
builder = builder_factory(
    "faiss",
    index_type="HNSW",  # 基于图的近似搜索
    metric="IP",  # 内积，适合归一化向量
    m=32,  # 每个节点的最大邻居数
    efSearch=64  # 搜索时的候选节点数
)

# 从嵌入向量构建索引并查询
with builder.from_embeddings(item_embeddings) as indexer:
    ids, distances = indexer.query(user_embeddings, top_k=20)
    indexer.save("faiss_hnsw.index")  # 保存索引

3. Milvus索引（分布式，适合生产环境）

Milvus是一个云原生向量数据库，支持分布式部署，适合大规模生产环境：

# 创建Milvus HNSW索引构建器
builder = builder_factory(
    "milvus",
    d=64,  # 向量维度
    index_type="HNSW",  # 基于图的索引
    metric="COSINE",  # 余弦相似度
    m=32,  # 每个节点的最大邻居数
    ef=64  # 搜索时的候选节点数
)

# 从嵌入向量构建索引并查询
with builder.from_embeddings(item_embeddings) as indexer:
    ids, distances = indexer.query(user_embeddings, top_k=10)

注意：Milvus需要单独启动服务，且当前封装主要用于实验，生产环境建议直接使用Milvus客户端管理collection。

索引后端选型建议

不同的索引后端各有特点，选择时需要根据实际场景权衡：

• 快速原型/小数据集：选择Annoy，安装简单，内存友好
• 中大规模离线计算：选择FAISS，性能优异，支持多种索引类型
• 大规模在线服务：选择Milvus，支持分布式部署和动态更新

完整部署示例：从训练到服务

下面我们以一个完整的示例，展示如何使用Torch-RecHub实现从模型训练到部署的全流程。

1. 训练模型并导出ONNX

# 假设我们已经训练好了一个DSSM模型
from torch_rechub.trainers import MatchTrainer

# 导出用户塔和物品塔
trainer.export_onnx("user_tower.onnx", mode="user")
trainer.export_onnx("item_tower.onnx", mode="item")

2. 量化ONNX模型

from torch_rechub.utils.quantization import quantize_model

# 量化用户塔
quantize_model("user_tower.onnx", "user_tower_int8.onnx", mode="int8")

# 量化物品塔
quantize_model("item_tower.onnx", "item_tower_int8.onnx", mode="int8")

3. 构建向量索引

from torch_rechub.serving import builder_factory

# 加载物品嵌入向量（假设已通过物品塔模型生成）
item_embeddings = ...  # 形状为 (n_items, embedding_dim)

# 使用FAISS构建向量索引
builder = builder_factory("faiss", index_type="HNSW", metric="IP", m=32)
with builder.from_embeddings(item_embeddings) as indexer:
    indexer.save("item_index.faiss")

4. 在线服务示例

import onnxruntime as ort
from torch_rechub.serving import builder_factory

# 加载用户塔ONNX模型
user_session = ort.InferenceSession("user_tower_int8.onnx")

# 加载物品向量索引
builder = builder_factory("faiss", index_type="HNSW", metric="IP")
indexer = builder.from_index_file("item_index.faiss")

def recommend(user_features, top_k=10):
    # 生成用户嵌入向量
    user_embedding = user_session.run(None, {"input": user_features})[0]
    
    # 查询相似物品
    item_ids, scores = indexer.query(user_embedding, top_k=top_k)
    
    return item_ids, scores

总结与最佳实践

通过本教程，我们学习了如何使用Torch-RecHub实现推荐模型的ONNX导出、量化优化和向量索引构建，从而完成从训练到生产环境的一键部署。以下是一些最佳实践建议：

1. 模型导出：优先使用动态batch size，提高服务灵活性
2. 量化优化：CPU环境推荐INT8量化，GPU环境推荐FP16转换
3. 向量索引：根据数据规模和应用场景选择合适的索引后端
4. 性能验证：使用提供的基准测试脚本验证导出和量化后的模型性能
5. 线上监控：部署后持续监控模型性能和推荐效果，必要时进行优化调整

Torch-RecHub提供了完整的推荐系统部署工具链，帮助开发者快速将研究成果转化为实际业务价值。更多详细信息，请参考官方文档和示例代码。

希望本教程能帮助你顺利实现推荐系统的部署，如有任何问题，欢迎在项目GitHub仓库提交issue或参与讨论。

【免费下载链接】torch-rechub A Lighting Pytorch Framework for Recommendation Models, Easy-to-use and Easy-to-extend. 项目地址: https://gitcode.com/gh_mirrors/to/torch-rechub