4、模型部署与推理优化：模型转换、量化技术与内存管理

模型训练出来只是第一步，真正让它跑在芯片上，才是见真章的时候。我这些年见过太多团队，模型精度刷得漂亮，一到部署就翻车。说白了，部署优化就是要把算法工程师的「艺术品」，变成硬件能高效执行的「工业品」。

这一章，我们聊聊三个核心环节：模型转换、量化技术、内存与数据搬运优化。这三件事环环相扣，任何一个环节出问题，推理性能都会大打折扣。

4.1 模型转换：从训练框架到推理引擎

你想想看，PyTorch 训练好的模型，能直接扔到 TensorRT 上跑吗？不行。中间需要一座桥——ONNX（Open Neural Network Exchange）。

ONNX 说白了就是一种中间表示格式。它不依赖任何训练框架，也不绑定任何推理引擎。我习惯把 ONNX 比作「通用语言」：你用 PyTorch 写，我用 TensorFlow 写，最后都转成 ONNX，大家就能互相理解了。

4.1.1 导出 ONNX 的常见坑

我在项目中遇到过最头疼的问题，就是动态尺寸。很多模型输入是固定尺寸（比如 224x224），但实际部署时可能需要处理不同分辨率的图片。导出 ONNX 时如果没设置动态轴，推理引擎就傻眼了。

# PyTorch 导出 ONNX 示例
import torch
import torch.onnx

model = torch.load('model.pth')
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', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size'}
    },
    opset_version=11
)

4.1.2 从 ONNX 到 TensorRT

ONNX 只是中间态，真正要跑在 NVIDIA 的 GPU 上，还得靠 TensorRT。TensorRT 会做三件事：层融合、精度校准、内核自动调优。

层融合是什么意思？举个例子，Conv + BN + ReLU 这三个操作，在训练时是分开的。但推理时，TensorRT 会把它们合并成一个 kernel。这样一来，显存访问次数减少，计算效率提升。我见过一个模型，光靠层融合就提速了 30%。

# TensorRT 构建引擎示例
import tensorrt as trt

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)

# 解析 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
serialized_engine = builder.build_serialized_network(network, config)

4.2 量化技术：INT8 与 FP16 的取舍

量化，说白了就是用更少的 bit 来表示模型参数和激活值。FP32 是 32 位浮点，FP16 是 16 位，INT8 是 8 位整数。位数越少，模型越小，推理越快，但精度可能会掉。

4.2.1 FP16：低风险，高回报

我个人建议，如果硬件支持 FP16（现在大部分 GPU 都支持），优先用 FP16。为什么？因为 FP16 的精度损失几乎可以忽略不计，但显存占用直接减半，吞吐量能翻倍。

我记得有一次优化一个 BERT 模型，FP32 下显存占用 12GB，换成 FP16 后降到 6GB，batch size 直接翻倍，推理速度提升了 1.8 倍。精度只掉了 0.1 个点，完全可接受。

4.2.2 INT8：需要校准，但效果惊人

INT8 才是真正的「大杀器」。模型大小压缩到 FP32 的 1/4，推理速度能提升 3-4 倍。但代价是：你需要做校准。

校准是什么意思？就是找一批有代表性的数据，跑一遍模型，统计出每层激活值的分布。然后根据这个分布，确定 INT8 的缩放因子。校准集选得好不好，直接决定量化后的精度。

# TensorRT INT8 校准器示例
import tensorrt as trt
import pycuda.driver as cuda

class Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_data, batch_size):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calibration_data = calibration_data
        self.batch_size = batch_size
        self.current_index = 0
        
        # 分配 GPU 内存
        self.device_input = cuda.mem_alloc(
            self.calibration_data[0].nbytes * self.batch_size
        )
    
    def get_batch(self, names):
        if self.current_index + self.batch_size > len(self.calibration_data):
            return None
        
        batch = self.calibration_data[
            self.current_index:self.current_index + self.batch_size
        ]
        self.current_index += self.batch_size
        
        cuda.memcpy_htod(self.device_input, batch)
        return [int(self.device_input)]
    
    def get_batch_size(self):
        return self.batch_size
    
    def read_calibration_cache(self):
        return None
    
    def write_calibration_cache(self, cache):
        with open('calibration.cache', 'wb') as f:
            f.write(cache)

4.2.3 量化精度对比

精度格式	模型大小	推理速度	精度损失	适用场景
FP32	基准	基准	无	精度要求极高的场景
FP16	50%	1.5x - 2x	几乎无	大部分推理场景
INT8	25%	3x - 4x	0.5% - 2%	对延迟敏感的场景

4.3 内存管理与数据搬运优化

模型转换和量化搞定了，但推理速度还是上不去？问题很可能出在内存上。GPU 算力再强，如果数据搬不过来，也是白搭。

4.3.1 显存分配策略

我见过很多新手，每次推理都重新分配显存。这其实非常慢。显存分配本身就有开销，频繁分配还会导致碎片化。

正确的做法是：一次性分配好所有需要的显存，然后复用。TensorRT 的 IExecutionContext 就支持这种方式，你可以在创建引擎时指定 workspace 大小，之后所有中间结果都在这个空间里复用。

4.3.2 数据搬运：CPU 与 GPU 之间的瓶颈

CPU 到 GPU 的数据搬运，是通过 PCIe 总线完成的。PCIe 4.0 x16 的理论带宽是 32GB/s，听起来很快对吧？但跟 GPU 内部 HBM 的 1TB/s 带宽比起来，就是小水管。

所以，优化的核心思路是：尽量减少 CPU 和 GPU 之间的数据搬运次数。

• 批量处理：把多个输入拼成一个 batch，一次搬运，一次推理。
• 异步传输：使用 CUDA Stream，让数据搬运和 kernel 执行重叠起来。
• 固定内存：使用 cudaHostAlloc 分配固定内存，能提升 PCIe 传输速度。

# 异步传输示例
import torch
import torch.cuda as cuda

# 创建 stream
stream = cuda.Stream()

# 分配固定内存
input_data = torch.empty((batch_size, 3, 224, 224), 
                          dtype=torch.float32, 
                          device='cpu').pin_memory()

# 异步拷贝到 GPU
with cuda.stream(stream):
    gpu_input = input_data.to('cuda', non_blocking=True)
    output = model(gpu_input)
    gpu_output = output.to('cpu', non_blocking=True)

# 等待完成
stream.synchronize()

4.3.3 内存布局与数据对齐

嗯，这里有个容易被忽略的点：数据对齐。GPU 对内存访问有对齐要求，通常是 128 字节或 256 字节。如果你的数据没对齐，GPU 会多读几次，性能直接下降。

我遇到过最离谱的一次，只是把输入图像的宽度从 224 改成 256，推理速度就快了 15%。为什么？因为 256 是 128 的倍数，数据对齐了。

4.4 本章小结

模型部署与推理优化，说白了就是三件事：

• 模型转换：用 ONNX 做桥梁，用 TensorRT 做引擎。注意动态尺寸和 opset 版本。
• 量化技术：FP16 是低风险首选，INT8 是性能利器但需要校准。校准集一定要有代表性。
• 内存与数据搬运：减少搬运次数，使用异步传输，注意数据对齐。这三条做好了，推理性能能再上一个台阶。

我做了这么多年芯片 SDK，最大的体会是：部署优化没有银弹。每个模型、每块芯片都有自己的脾气。你需要理解底层原理，然后针对性地调优。希望这一章的内容，能帮你少走一些弯路。

---