📚 课程目录第3章 CPU推理优化

3. CPU推理优化

任务目标

将训练好的PyTorch模型导出为ONNX格式，使用ONNX Runtime在CPU上完成推理，并与原始PyTorch模型的推理速度和内存占用进行对比，验证ONNX Runtime在CPU场景下的优化效果。

环境依赖

Python 3.8+
PyTorch 1.12+
ONNX 1.12+
ONNX Runtime（CPU版本）
psutil（内存测量）

安装命令：pip install torch onnx onnxruntime psutil

完整实现代码

import torch
import torch.nn as nn
import onnx
import onnxruntime as ort
import time
import psutil
import os

# 1. 定义并初始化一个简单的PyTorch模型（用于演示）
class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleMLP()
model.eval()

# 2. 构造示例输入并导出ONNX
dummy_input = torch.randn(1, 784)
onnx_path = "simple_mlp.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    input_names=["input"],
    output_names=["output"],
    opset_version=13,
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

# 验证ONNX模型完整性
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("ONNX模型导出并验证通过")

# 3. 创建ONNX Runtime推理会话（CPU）
ort_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])

# 4. 准备测试输入（固定随机种子保证可复现）
torch.manual_seed(42)
test_input = torch.randn(8, 784)
ort_input = {ort_session.get_inputs()[0].name: test_input.numpy()}

# 5. 推理速度对比（预热后各测50次）
def measure_pytorch_inference(model, input_tensor, iterations=50):
    with torch.no_grad():
        for _ in range(5):  # 预热
            model(input_tensor)
        start = time.perf_counter()
        for _ in range(iterations):
            model(input_tensor)
        elapsed = time.perf_counter() - start
    return elapsed / iterations

def measure_ort_inference(session, input_dict, iterations=50):
    for _ in range(5):  # 预热
        session.run(None, input_dict)
    start = time.perf_counter()
    for _ in range(iterations):
        session.run(None, input_dict)
    elapsed = time.perf_counter() - start
    return elapsed / iterations

pytorch_avg_time = measure_pytorch_inference(model, test_input)
ort_avg_time = measure_ort_inference(ort_session, ort_input)

# 6. 内存占用对比（测量推理过程中的峰值内存增量）
def measure_memory_increase(func, *args):
    process = psutil.Process(os.getpid())
    mem_before = process.memory_info().rss / 1024 / 1024  # MB
    func(*args)
    mem_after = process.memory_info().rss / 1024 / 1024
    return mem_after - mem_before

# PyTorch内存增量（执行10次推理取平均）
pytorch_mem = measure_memory_increase(
    lambda: [model(test_input) for _ in range(10)]
) / 10

# ONNX Runtime内存增量
ort_mem = measure_memory_increase(
    lambda: [ort_session.run(None, ort_input) for _ in range(10)]
) / 10

# 7. 输出对比结果
print(f"\n=== CPU推理性能对比 ===")
print(f"PyTorch平均推理时间: {pytorch_avg_time*1000:.2f} ms")
print(f"ONNX Runtime平均推理时间: {ort_avg_time*1000:.2f} ms")
print(f"速度提升: {(pytorch_avg_time/ort_avg_time):.2f}x")
print(f"\nPyTorch平均内存增量: {pytorch_mem:.2f} MB")
print(f"ONNX Runtime平均内存增量: {ort_mem:.2f} MB")
print(f"内存减少: {max(0, (1 - ort_mem/pytorch_mem)*100):.1f}%")

# 8. 验证输出一致性
with torch.no_grad():
    pytorch_output = model(test_input).numpy()
ort_output = ort_session.run(None, ort_input)[0]
max_diff = np.abs(pytorch_output - ort_output).max()
print(f"\n输出最大差异: {max_diff:.2e}")
print(f"结果一致性: {'通过' if max_diff < 1e-4 else '失败'}")

运行方式

将上述代码保存为 cpu_optimization.py；
在终端执行 python cpu_optimization.py；
脚本会自动完成模型导出、推理测试和结果对比。

验证方法与判定标准

导出验证：控制台输出“ONNX模型导出并验证通过”，说明ONNX格式正确；
性能对比：观察输出的平均推理时间，ONNX Runtime通常比PyTorch快1.5～3倍（具体取决于CPU型号）；
内存对比：ONNX Runtime的内存增量通常低于PyTorch；
一致性验证：输出最大差异小于1e-4时显示“通过”，确保转换未破坏模型精度。

常见问题

ONNX导出失败：确认PyTorch版本≥1.12，且模型不含自定义算子；
速度提升不明显：模型过小时优化空间有限，可增大输入batch或使用更深网络验证；
内存测量波动：建议关闭其他程序后测试，或增加推理次数取平均值。