4. 模型推理优化：模型量化与剪枝、ONNX Runtime与TensorRT部署、GPU推理加速与批处理

模型训练完了，上线一跑，延迟高得吓人。这事我遇到过不止一次。

说实话，很多团队把精力都花在训练阶段，觉得模型精度高就万事大吉。但到了生产环境，推理速度才是真正的瓶颈。你想想看，推荐系统要求毫秒级响应，模型再准，跑不动也是白搭。

这一章，咱们就聊聊怎么让模型跑得快。核心就三件事：模型瘦身、推理引擎加速、GPU榨干性能。

4.1 模型量化：把浮点数变成整数

模型量化，说白了就是把模型里的参数从32位浮点数（FP32）压缩成16位（FP16）甚至8位整数（INT8）。

为什么能这么做？因为推荐模型里很多参数对精度没那么敏感。我做过一个实验，把DIN模型从FP32量化到INT8，AUC只掉了0.1%，但推理速度提升了3倍。

量化带来的收益：

  • 模型体积缩小：FP32→INT8，体积直接缩到1/4
  • 内存带宽降低：加载模型更快，缓存命中率更高
  • 计算速度提升：INT8的乘加运算比FP32快2-4倍

量化有两种方式：

  • 训练后量化（PTQ）：训练完直接转，简单粗暴。我建议先用这个试试，大部分场景够用。
  • 量化感知训练（QAT）：训练时就模拟量化误差，精度更高。如果PTQ掉点严重，再上这个。

我的经验：做量化时，先对模型做一次校准（calibration）。拿一批真实数据跑一遍，统计出每层激活值的分布范围。这样量化后的精度损失最小。我曾经跳过这步，结果线上AUC掉了0.5%，被老板叫去喝茶了。

4.2 模型剪枝：砍掉没用的参数

模型剪枝，就是去掉那些不重要的权重。你想想看，一个深度推荐模型可能有上亿参数，但真正起作用的可能只有20%。

剪枝分两种：

  • 结构化剪枝：直接砍掉整个神经元或通道。好处是能真正减少计算量，坏处是精度损失大。
  • 非结构化剪枝：只把权重值小的参数置零。精度损失小，但需要特殊硬件支持才能加速。

我个人习惯用结构化剪枝。虽然调参麻烦点，但效果立竿见影。

注意：剪枝不是一次搞定的。我建议用「迭代剪枝」策略：剪掉5%→微调→再剪5%→再微调。这样模型有足够时间适应，精度不会崩。

剪枝后的模型，配合量化一起用，效果更佳。我做过一个案例：先剪掉40%的参数，再量化到INT8，模型体积缩到原来的1/10，推理速度提升5倍，AUC只掉了0.3%。

4.3 ONNX Runtime与TensorRT部署

模型训练好之后，怎么部署到生产环境？直接用PyTorch或TensorFlow的推理接口？说实话，太慢了。

这里我推荐两个工具：ONNX Runtime 和 TensorRT。

4.3.1 ONNX Runtime：通用推理引擎

ONNX（Open Neural Network Exchange）是一个开放的模型格式。你训练好的模型，转成ONNX格式后，就能用ONNX Runtime来跑。

好处是什么？

  • 跨框架：不管你是PyTorch、TensorFlow还是Keras，都能转成ONNX
  • 跨平台：Windows、Linux、ARM都能跑
  • 自动优化：ONNX Runtime会自动做图优化、算子融合
# 把PyTorch模型转成ONNX
import torch
import torch.onnx

model = YourModel()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "model.onnx",
                  input_names=['input'],
                  output_names=['output'],
                  dynamic_axes={'input': {0: 'batch_size'},
                                'output': {0: 'batch_size'}})

避坑指南：导出ONNX时，一定要设置dynamic_axes。否则你的模型只能跑固定batch size，线上流量波动时就傻眼了。我曾经因为这个，线上服务半夜挂了，被运维同事骂了一顿。

4.3.2 TensorRT：NVIDIA的杀手锏

如果你用的是NVIDIA GPU，那TensorRT就是你的首选。它比ONNX Runtime还能再快2-3倍。

TensorRT的核心优化：

  • 算子融合：把多个小算子合并成一个，减少kernel launch开销
  • 精度校准：自动选择FP16或INT8，在精度和速度间找平衡
  • 内存优化：复用中间结果，减少显存占用
# TensorRT推理示例
import tensorrt as trt
import pycuda.driver as cuda

# 创建TensorRT引擎
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network()
parser = trt.OnnxParser(network, logger)

# 解析ONNX模型
with open("model.onnx", "rb") as f:
    parser.parse(f.read())

# 构建引擎
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
engine = builder.build_serialized_network(network, config)

我的建议：先用ONNX Runtime做快速部署，等系统稳定了，再换成TensorRT做极致优化。别一上来就上TensorRT，调试起来太痛苦了。

4.4 GPU推理加速与批处理

GPU的并行计算能力很强，但如果你不会用，它就是个摆设。

4.4.1 动态批处理

推荐系统的请求是实时到达的，每个请求的batch size=1。但GPU处理单个请求的效率很低，因为kernel launch的开销是固定的。

解决方案：动态批处理。把多个请求攒起来，凑成一个batch再送进GPU。

# 动态批处理伪代码
class BatchProcessor:
    def __init__(self, max_batch_size=64, max_wait_ms=10):
        self.queue = []
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
    
    def process(self, request):
        self.queue.append(request)
        if len(self.queue) >= self.max_batch_size:
            return self.flush()
        # 等待更多请求或超时
        time.sleep(self.max_wait_ms / len(self.queue))
        return self.flush()
    
    def flush(self):
        batch = self.queue[:self.max_batch_size]
        self.queue = self.queue[self.max_batch_size:]
        return model.inference(batch)

注意：批处理不是越大越好。batch size太大，单个请求的等待时间会变长。我一般控制在16-64之间，延迟和吞吐量能取得平衡。

4.4.2 GPU内存管理

GPU显存是稀缺资源。推荐模型动辄几百MB，多个模型同时部署时，显存很容易爆。

我的做法：

  • 模型分片：把大模型拆成多个小模型，按需加载
  • 显存池化：预分配显存，避免频繁申请释放
  • 混合精度推理：用FP16代替FP32，显存直接减半

4.4.3 异步推理

别让CPU等着GPU算完。用异步推理，CPU可以同时处理其他请求。

# 异步推理示例
import asyncio
import torch

async def async_inference(model, input_tensor):
    with torch.no_grad():
        # 把计算放到GPU上，不阻塞CPU
        future = torch.jit.fork(model, input_tensor)
        result = await torch.jit.wait(future)
    return result

# 并发处理多个请求
async def handle_requests(requests):
    tasks = [async_inference(model, req) for req in requests]
    results = await asyncio.gather(*tasks)
    return results

我的经验：异步推理配合动态批处理，能把GPU利用率从30%提升到80%以上。但要注意，异步逻辑容易出bug，一定要做好超时和错误处理。

4.5 实战：一个完整的优化流程

说了这么多，咱们串起来看看一个完整的优化流程：

  1. 模型瘦身：先做结构化剪枝，砍掉30%的参数。然后做训练后量化，转成INT8。
  2. 格式转换：把PyTorch模型导出为ONNX格式，设置动态batch。
  3. 推理引擎：先用ONNX Runtime部署，验证精度和速度。然后换成TensorRT，再做一轮优化。
  4. GPU加速：实现动态批处理，batch size设为32。开启异步推理，配合显存池化。
  5. 压测验证：用真实流量压测，确保P99延迟在50ms以内。

这套流程我用了很多次，效果很稳定。有一次，一个原本需要4张V100才能扛住的模型，优化后只用1张T4就搞定了。老板看了成本报表，笑得合不拢嘴。

最后说一句：模型推理优化没有银弹。量化、剪枝、TensorRT、批处理，每个手段都有适用场景。我的建议是：先做性能分析，找到真正的瓶颈，再对症下药。别一上来就全上，容易把自己搞晕。

← 上一章 📖 返回目录 下一章 →