11. 模型部署与推理优化

ML 模型生产化部署指南：ONNX 跨平台转换、TorchServe 服务化、量化蒸馏压缩、TensorRT 加速与边缘端部署

Leeting Yan 2026-08-13 7 分钟阅读 3272 字

训练好的模型只有部署到生产环境才能创造价值。本文覆盖模型格式转换、服务化架构、推理加速（量化、剪枝、蒸馏）与边缘端部署的完整链路。

1. 模型持久化与格式

1.1 PyTorch 模型保存

import torch

# 保存完整模型（不推荐，依赖环境）
torch.save(model, 'model.pt')
model = torch.load('model.pt')

# 保存状态字典（推荐，跨版本兼容）
torch.save(model.state_dict(), 'model_weights.pth')

# 加载
model = MyModel()
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

# 保存训练状态（含优化器、epoch）
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, 'checkpoint.tar')

1.2 Pickle vs Safetensors

# Safetensors：安全、快速、无代码执行风险
from safetensors.torch import save_file, load_file

save_file(model.state_dict(), 'model.safetensors')
state_dict = load_file('model.safetensors')

2. ONNX 跨平台部署

2.1 PyTorch → ONNX

ONNX (Open Neural Network Exchange) 是跨框架的中间表示格式。

import torch

# 导出 ONNX
model.eval()
dummy_input = torch.randn(1, 3, 224, 224).to('cuda')

torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    export_params=True,           # 存储训练好的参数
    opset_version=17,             # ONNX 算子集版本
    do_constant_folding=True,     # 常量折叠优化
    input_names=['input'],        # 输入名
    output_names=['output'],      # 输出名
    dynamic_axes={                # 动态维度（batch size）
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

2.2 ONNX Runtime 推理

import onnxruntime as ort
import numpy as np

# 创建 InferenceSession
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
session = ort.InferenceSession('model.onnx', providers=providers)

# 查看输入输出信息
for inp in session.get_inputs():
    print(f"Input: {inp.name}, Shape: {inp.shape}, Type: {inp.type}")

# 推理
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

x = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run([output_name], {input_name: x})

# 批量推理
x_batch = np.random.randn(32, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {input_name: x_batch})

2.3 ONNX 优化

import onnx
from onnxruntime.tools import optimizer

# 加载并检查
model = onnx.load('model.onnx')
onnx.checker.check_model(model)

# 图优化
optimized_model = optimizer.optimize_model(
    'model.onnx',
    model_type='bert',  # 或 'bert', 'gpt2', 'bart'
    num_heads=12,
    hidden_size=768,
    use_gpu=True
)
optimized_model.convert_model_float32_to_float16()
optimized_model.save_model_to_file('model_optimized.onnx')

3. TorchServe 服务化

3.1 搭建模型服务

# model_handler.py
from ts.torch_handler.base_handler import BaseHandler
import torch
import json
import os

class CustomHandler(BaseHandler):
    def initialize(self, context):
        self.manifest = context.manifest
        properties = context.system_properties
        model_dir = properties.get("model_dir")
        
        # 加载模型
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = torch.jit.load(os.path.join(model_dir, "model.pt"))
        self.model.to(self.device)
        self.model.eval()
        
        # 加载预处理器
        self.transform = ...
    
    def preprocess(self, data):
        images = []
        for row in data:
            image = row.get("data") or row.get("body")
            image = Image.open(io.BytesIO(image))
            image = self.transform(image)
            images.append(image)
        return torch.stack(images).to(self.device)
    
    def inference(self, data):
        with torch.no_grad():
            return self.model(data)
    
    def postprocess(self, inference_output):
        probs = torch.softmax(inference_output, dim=1)
        return json.dumps([{"class": p.argmax().item(), 
                           "confidence": p.max().item()} for p in probs])

3.2 打包与部署

# 1. 创建 MAR 模型归档
torch-model-archiver --model-name my_model \
  --version 1.0 \
  --model-file model.py \
  --serialized-file model.pth \
  --handler custom_handler.py \
  --export-path model_store

# 2. 启动 TorchServe
torchserve --start \
  --model-store model_store \
  --models my_model=my_model.mar \
  --ts-config config.properties

# 3. 测试推理
curl -X POST http://localhost:8080/predictions/my_model \
  -T sample.jpg

# 4. 管理 API
# 查看模型列表
curl http://localhost:8081/models
# 弹性伸缩 workers
curl -X PUT "http://localhost:8081/models/my_model?min_worker=2&max_worker=4"

3.3 配置优化

# config.properties
inference_address=http://0.0.0.0:8080
management_address=http://0.0.0.0:8081
metrics_address=http://0.0.0.0:8082

number_of_netty_threads=32
job_queue_size=1000
async_logging=true

# 默认 workers
default_workers_per_model=2

4. 模型量化

量化将 FP32 权重降低为 INT8/FP16，减少内存占用和计算量。

4.1 PyTorch 动态量化

import torch.quantization

# 动态量化（最简单，适用于 LSTM/Transformer）
model_int8 = torch.quantization.quantize_dynamic(
    model,
    {nn.Linear, nn.LSTM},  # 要量化的层
    dtype=torch.qint8
)

# 对比大小
print(f"FP32 大小: {model.get_size()}")
print(f"INT8 大小: {model_int8.get_size()}")

4.2 静态量化（PTQ）

# 准备模型
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model, inplace=True)

# 用校准数据收集统计信息
with torch.no_grad():
    for x, _ in calib_loader:
        model(x)

# 转换
torch.quantization.convert(model, inplace=True)

4.3 量化感知训练 (QAT)

# 在训练中模拟量化效果，精度更高
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)

# 正常训练几个 epoch
for epoch in range(3):
    train(model, ...)

# 冻结 BN 统计并转换
model.apply(torch.quantization.disable_observer)
torch.quantization.convert(model, inplace=True)

4.4 ONNX 量化

from onnxruntime.quantization import quantize_dynamic, QuantType

quantize_dynamic(
    model_input='model.onnx',
    model_output='model_int8.onnx',
    weight_type=QuantType.QInt8
)

量化方式	方法	精度损失	加速比	适用
Post-Training	校准后静态量化	中	2-4x	CNN、NLP
Quantization-Aware	训练中模拟	低	2-4x	对精度敏感
Dynamic	运行时动态	中高	1.5-2x	快速上手

5. 知识蒸馏 (Knowledge Distillation)

让小模型（学生）学习大模型（教师）的输出分布。

class DistillationLoss(nn.Module):
    def __init__(self, temperature=4.0, alpha=0.7):
        super().__init__()
        self.T = temperature
        self.alpha = alpha
        self.ce = nn.CrossEntropyLoss()
        self.kl = nn.KLDivLoss(reduction='batchmean')
    
    def forward(self, student_logits, teacher_logits, labels):
        # Hard loss（学生自己学标签）
        hard_loss = self.ce(student_logits, labels)
        
        # Soft loss（学生学习教师"软标签"）
        student_soft = F.log_softmax(student_logits / self.T, dim=1)
        teacher_soft = F.softmax(teacher_logits / self.T, dim=1)
        soft_loss = self.kl(student_soft, teacher_soft) * (self.T ** 2)
        
        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

# 训练循环
teacher.eval()
for epoch in range(epochs):
    student.train()
    for x, y in train_loader:
        optimizer.zero_grad()
        
        with torch.no_grad():
            teacher_logits = teacher(x)
        
        student_logits = student(x)
        loss = distillation_loss(student_logits, teacher_logits, y)
        loss.backward()
        optimizer.step()

6. TensorRT 加速

6.1 ONNX → TensorRT

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

# 1. 创建 builder 和 logger
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, logger)

# 2. 解析 ONNX
with open('model.onnx', 'rb') as f:
    parser.parse(f.read())

# 3. 配置 builder
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30  # 1GB
config.set_flag(trt.BuilderFlag.FP16)

# 4. 构建引擎
engine = builder.build_engine(network, config)

# 5. 序列化保存
with open('model.trt', 'wb') as f:
    f.write(engine.serialize())

6.2 TensorRT 推理

# 反序列化引擎
with open('model.trt', 'rb') as f:
    runtime = trt.Runtime(logger)
    engine = runtime.deserialize_cuda_engine(f.read())

context = engine.create_execution_context()

# 分配 GPU 内存
h_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)
h_output = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)), dtype=np.float32)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
stream = cuda.Stream()

# 推理
cuda.memcpy_htod_async(d_input, h_input, stream)
context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
cuda.memcpy_dtoh_async(h_output, d_output, stream)
stream.synchronize()

6.3 torch2trt

from torch2trt import torch2trt

model = model.eval().cuda()
x = torch.ones((1, 3, 224, 224)).cuda()

# 转换 FP16
model_trt = torch2trt(model, [x], fp16_mode=True)

# 推理
y_trt = model_trt(x)

7. 模型剪枝

移除不重要权重，减少参数量和计算量。

import torch.nn.utils.prune as prune

# 非结构化剪枝（随机）
prune.random_unstructured(module, name='weight', amount=0.3)

# 结构化剪枝（L1 范数）
prune.ln_structured(module, name='weight', amount=0.3, n=1, dim=0)

# 全局剪枝
parameters_to_prune = [
    (model.conv1, 'weight'),
    (model.conv2, 'weight'),
    (model.fc1, 'weight'),
]
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.2
)

8. 边缘端部署

8.1 移动端：Core ML / TFLite

# PyTorch → ONNX → Core ML
import coremltools as ct

# 从 ONNX 转换
mlmodel = ct.converters.onnx.convert(
    model='model.onnx',
    minimum_ios_deployment_target='13'
)
mlmodel.save('model.mlmodel')

# PyTorch → TFLite (通过 ONNX 中转)
# torch → ONNX → TensorFlow → TFLite

8.2 树莓派 / Jetson

# Jetson Nano/Xavier 用 JetPack + TensorRT
# ARM 架构需交叉编译或使用 conda-forge

# 安装 PyTorch for ARM
wget https://nvidia.box.com/some_path/torch-2.0.0-cp38-cp38-linux_aarch64.whl
pip install torch-2.0.0-cp38-cp38-linux_aarch64.whl

9. 推理服务架构

9.1 Triton Inference Server

# 目录结构
model_repository/
  └── my_model/
      ├── 1/
      │   └── model.onnx
      └── config.pbtxt

# config.pbtxt
name: "my_model"
platform: "onnxruntime_onnx"
max_batch_size: 32
input [
  {
    name: "input"
    data_type: TYPE_FP32
    dims: [3, 224, 224]
  }
]
output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [1000]
  }
]
instance_group [
  {
    count: 2
    kind: KIND_GPU
  }
]

# 启动
docker run --gpus all -p 8000:8000 -p 8001:8001 -p 8002:8002 \
  -v $(pwd)/model_repository:/models \
  nvcr.io/nvidia/tritonserver:23.10-py3 \
  tritonserver --model-repository=/models

9.2 性能基准测试

import time

# 预热
for _ in range(10):
    _ = model(dummy_input)

torch.cuda.synchronize()

# 测试
start = time.time()
for _ in range(100):
    output = model(batch_input)
torch.cuda.synchronize()
end = time.time()

throughput = (100 * batch_size) / (end - start)
latency_ms = (end - start) / 100 * 1000
print(f"吞吐: {throughput:.1f} samples/s, 延迟: {latency_ms:.2f} ms")

10. 模型版本管理与追踪

生产环境的模型迭代需要严格的版本控制。MLflow 和 Weights & Biases 是业界主流方案：

import mlflow

with mlflow.start_run():
    # 记录超参数
    mlflow.log_params({"lr": 0.001, "batch_size": 32, "epochs": 50})
    
    # 记录指标
    mlflow.log_metrics({"val_acc": 0.947, "f1": 0.932})
    
    # 记录模型
    mlflow.pytorch.log_model(model, "model")
    
    # 注册到 Model Registry
    model_version = mlflow.register_model(f"runs:/{mlflow.active_run().info.run_id}/model", "resnet50-classifier")

# 生产加载指定版本
model_uri = f"models:/{model_name}/{stage}"  # stage 可以是 Production / Staging / Archived
model = mlflow.pytorch.load_model(model_uri)

版本管理的关键实践：

每个训练 run 生成唯一版本号，不可变
通过 stage 标签（Staging → Production → Archived）管理生命周期
部署时通过版本号回滚，确保可复现性

11. A/B 测试与阴影部署

除了简单的流量切分，生产模型部署还需支持渐进式验证：

A/B 测试：不同模型处理不同用户，对比业务指标（转化率、点击率、留存）

class ModelRouter:
    def __init__(self, model_a, model_b, split=0.5):
        self.model_a = model_a
        self.model_b = model_b
        self.split = split
    
    def predict(self, x, user_id):
        if hash(user_id) % 100 < self.split * 100:
            return self.model_a(x), 'A'
        return self.model_b(x), 'B'

阴影部署（Shadow Deployment）：新模型并行处理生产流量，返回旧模型结果，仅收集新模型的延迟、输出分布用于验证，零风险。

def shadow_predict(request):
    # 主模型（旧版）返回结果
    primary_result = model_v1.predict(request)
    
    # 阴影模型（新版）异步执行，只记录不返回
    asyncio.create_task(log_shadow_result(model_v2.predict(request)))
    
    return primary_result

12. LLM 推理优化技术

大语言模型的推理优化是当前部署热点：

KV Cache 优化

自回归生成中，已计算的 Key/Value 缓存可复用，避免重复计算：

# 标准做法：每次传入完整上下文
for i in range(max_new_tokens):
    output = model(input_ids)  # 重复计算所有 history 的 K/V
    next_token = sampling(output[:, -1, :])
    input_ids = torch.cat([input_ids, next_token], dim=1)

# 优化后：维护 past_key_values 缓存
past_kv = None
for i in range(max_new_tokens):
    outputs = model(input_ids, past_key_values=past_kv, use_cache=True)
    past_kv = outputs.past_key_values  # (K, V) 追加而非重算
    next_token = sampling(outputs.logits[:, -1, :])
    input_ids = next_token  # 只传最新 token

KV Cache 将长序列生成的推理复杂度从 $O(n^2)$ 降至 $O(n)$，但会额外消耗显存（ batch_size × num_layers × num_heads × head_dim × seq_len × 2 bytes）。

多 GPU 并行策略

策略	原理	适用场景	通信开销
Pipeline Parallelism	层间拆分，不同 GPU 处理不同层	单节点多卡	较低（仅 activations）
Tensor Parallelism	层内拆分，同一层拆分到多卡	模型过大无法单卡加载	较高（每层的 all-reduce）
Sequence Parallelism	序列维度拆分	长序列 LLM	中

实际部署大型模型（如 70B）通常采用 张量并行（张量并行，TP=8）+ 流水线并行（PP=4） 的组合方案。

Batch 与 Throughput 优化

技术	效果	实现难度
Dynamic Batching	合并请求提升 GPU 利用率	低（vLLM / Triton 内置）
In-flight Batching	一个序列完成后立即替补新序列	中（Continuous Batching）
PagedAttention	非连续 KV Cache 存储，减少碎片	中（vLLM 核心创新）
Speculative Decoding	小模型草稿 + 大模型验证，翻倍吞吐	高

vLLM 通过 PagedAttention 机制将 LLM 推理吞吐提升了 2-4 倍，已成为 LLM 部署的事实标准。

13. 推理服务的自动扩缩容

Kubernetes 生态下，GPU 推理服务的弹性伸缩面临特殊挑战：

# KEDA ScaledObject 基于队列长度触发 GPU 扩缩容
apiVersion: keda.sh/v1alpha1
kind: ScaledObject
metadata:
  name: llm-inference-scaler
spec:
  scaleTargetRef:
    name: llm-inference
  triggers:
    - type: metrics-api
      metadata:
        url: "http://prometheus:9090/api/v1/query?query=pending_requests"
        valueLocation: "data.result.0.value.1"
        targetValue: "10"
  minReplicaCount: 1
  maxReplicaCount: 8

GPU 节点扩容通常需要 2-5 分钟，远慢于 CPU。推荐策略：

设置合理的 minReplicas，避免冷启动延迟拖垮体验
使用 Cluster Autoscaler 配合 GPU 节点池（预置镜像、提前拉取模型）
混合部署：轻量请求走 CPU 推理（ollama / llama.cpp），重载请求路由到 GPU 集群

14. 模型部署成本优化框架

策略	节省	权衡
模型量化（INT8/FP16）	50-75% 显存	轻微精度损失
量化的动态批处理	30-50% 成本	延迟波动
竞价实例（Spot）	60-90% 计算成本	可中断，需 checkpoint
共享 GPU（MIG / time-slicing）	多租户分摊	隔离与性能保障
边缘/端侧卸载	云成本降为零	模型大小受限

成本监控公式：

单次推理成本 = (GPU 小时单价 × 推理耗时) / batch_size
优化方向：提高 batch_size、降低单次耗时、切换到低价实例

15. 监控与 A/B 测试

# 模型性能监控指标
metrics = {
    'inference_latency_p99': np.percentile(latencies, 99),
    'throughput': samples_per_second,
    'gpu_utilization': gpustat_GPUUtilization(),
    'memory_usage_mb': torch.cuda.memory_allocated() / 1e6,
    'prediction_drift': ks_statistic(new_preds, baseline_preds),
    'data_drift': ks_statistic(new_data, training_data)
}

# A/B 测试
class ModelRouter:
    def __init__(self, model_a, model_b, split=0.5):
        self.model_a = model_a
        self.model_b = model_b
        self.split = split
    
    def predict(self, x, user_id):
        if hash(user_id) % 100 < self.split * 100:
            return self.model_a(x), 'A'
        return self.model_b(x), 'B'

总结

部署方案	工具	适用场景	延迟
本地推理	PyTorch/ONNX Runtime	开发调试	取决于硬件
GPU 加速	TensorRT	高吞吐 GPU 服务	< 10ms
服务化	TorchServe/Triton	多模型、高并发	< 50ms
移动端	Core ML / TFLite	iOS/Android App	< 100ms
边缘设备	OpenVINO / TensorRT	Jetson/树莓派	< 500ms

优化手段	效果	难度	适用
FP16	2x 加速，几乎无损	低	所有 GPU
INT8 量化	2-4x 加速，轻微损失	中	部署首选
知识蒸馏	小模型达到大模型 90%+	中	资源受限
剪枝	2-10x 压缩	中高	极端场景
ONNX Runtime	跨平台 2-3x	低	通用部署

部署最佳实践：

基线：先用 ONNX Runtime 快速部署，验证业务逻辑
优化：针对硬件选 TensorRT（NVIDIA）或 OpenVINO（Intel）
量化：PTQ 先做，精度不够再 QAT
服务化：TorchServe 快速搭建，Triton 管理复杂多模型场景
监控：预测分布漂移、输入数据漂移、延迟与错误率

继续阅读

探索更多技术文章

浏览归档，发现更多关于系统设计、工具链和工程实践的内容。

全部文章返回首页