TensorFlow模型转TensorRT加速

永远是夏天
2 阅读6分钟

你想了解如何用TensorRT对TensorFlow训练的模型进行加速，这是工业级部署中提升推理性能的核心手段，尤其适用于NVIDIA GPU场景（如服务器、边缘计算设备Jetson系列）。下面我会从TensorRT核心原理、TensorFlow模型转TensorRT流程、实战优化步骤、性能对比四个维度，详细讲解落地方案，所有代码均可直接复用。

一、先搞懂：TensorRT为什么能加速？

TensorRT是NVIDIA推出的高性能推理优化引擎，核心是对深度学习模型做编译级优化，相比原生TensorFlow/PyTorch推理，性能可提升2~10倍，核心优化手段包括：

  1. 算子融合：将多个连续的算子（如Conv+BN+ReLU）融合为单个算子，减少显存读写和内核调用开销；
  2. 精度校准：支持FP16（半精度）、INT8（整型）量化，在精度损失可控的前提下，提升计算速度、降低显存占用；
  3. 层/张量融合：消除冗余的激活函数、拼接操作，优化张量内存布局；
  4. 内核自动调优：根据GPU硬件特性（如算力架构、显存带宽），选择最优的计算内核；
  5. 动态张量显存：避免静态分配冗余显存，提升显存利用率。

二、核心流程：TensorFlow模型转TensorRT加速（实战步骤）

以TensorFlow训练的图像分类模型（MobileNetV2）为例，完整实现TensorRT加速，环境要求：

  • 硬件：NVIDIA GPU（算力≥6.0，如GTX 1080、T4、A100）；
  • 软件：CUDA 11.x、cuDNN 8.x、TensorRT 8.x、TensorFlow 2.x、onnxruntime-gpu。

步骤1：环境安装

# 1. 安装TensorRT（推荐用NVIDIA官方docker，避免环境冲突）
# 或手动安装：https://developer.nvidia.com/tensorrt
pip install tensorrt==8.6.1

# 2. 安装依赖库
pip install onnx onnxruntime-gpu tf2onnx  # 用于TF转ONNX

步骤2：准备TensorFlow模型（SavedModel格式）

TensorRT优先支持SavedModel格式，若你的模型是.h5格式，先转换：

import tensorflow as tf

# 1. 加载.h5模型（以之前的缺陷检测模型为例）
h5_model = tf.keras.models.load_model("best_model.h5")

# 2. 保存为SavedModel格式
saved_model_dir = "saved_model"
h5_model.save(saved_model_dir, save_format="tf")

步骤3：方式1：TensorFlow-TensorRT（TF-TRT）直接转换（最简单）

TF-TRT是TensorFlow内置的TensorRT集成模块，无需手动转格式，直接优化模型，适合快速落地：

import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt

# 1. 配置转换参数
converter = trt.TrtGraphConverterV2(
    input_saved_model_dir=saved_model_dir,
    # 精度模式：FP32/FP16/INT8（INT8需校准数据集）
    precision_mode=trt.TrtPrecisionMode.FP16,  
    max_workspace_size_bytes=1 << 30,  # 显存分配：1GB
    maximum_cached_engines=1  # 缓存引擎数量
)

# 2. 转换并保存优化后的模型
converter.convert()
trt_saved_model_dir = "trt_saved_model"
converter.save(trt_saved_model_dir)

# 3. 加载TRT优化后的模型推理
loaded_trt_model = tf.saved_model.load(trt_saved_model_dir)
infer = loaded_trt_model.signatures["serving_default"]

# 4. 测试推理（输入需为TensorFlow张量）
import numpy as np
# 构造测试输入：(1, 224, 224, 3)，归一化到0-1
test_input = tf.constant(np.random.rand(1, 224, 224, 3).astype(np.float32))
# 推理
result = infer(test_input)
# 解析输出（根据模型输出名调整，如"dense_1"）
pred = result[list(result.keys())[0]].numpy()
print(f"TRT推理结果：{np.argmax(pred)}")

步骤4：方式2：TF→ONNX→TensorRT（更灵活，精度更高）

若TF-TRT优化效果不佳，可通过ONNX中间格式转换，这是工业级常用方案：

子步骤4.1：TensorFlow模型转ONNX
import tf2onnx
import tensorflow as tf

# 加载SavedModel模型
model = tf.saved_model.load(saved_model_dir)
# 获取模型签名
signature = model.signatures["serving_default"]

# 转换为ONNX格式
onnx_model_path = "model.onnx"
spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)
output_path = tf2onnx.convert.from_keras_model(
    h5_model,  # 也可直接用h5模型
    input_signature=spec,
    opset=13,  # ONNX算子集版本，建议11+
    output_path=onnx_model_path
)
print(f"ONNX模型已保存至：{onnx_model_path}")
子步骤4.2：ONNX模型转TensorRT引擎（关键）
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit  # 自动初始化CUDA

# 1. 初始化TensorRT logger
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)

# 2. 解析ONNX模型
with open(onnx_model_path, "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))
        raise RuntimeError("解析ONNX模型失败")

# 3. 配置构建参数（核心优化）
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30  # 1GB显存
# 设置精度模式
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)  # FP16量化
# 若需INT8量化，需添加校准器（见下文）

# 4. 构建并保存TensorRT引擎
engine_path = "model.trt"
serialized_engine = builder.build_serialized_network(network, config)
with open(engine_path, "wb") as f:
    f.write(serialized_engine)
print(f"TensorRT引擎已保存至：{engine_path}")
子步骤4.3：加载TensorRT引擎推理（高性能）
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda

# 1. 加载TensorRT引擎
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
runtime = trt.Runtime(TRT_LOGGER)
with open("model.trt", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()

# 2. 分配显存（输入/输出）
def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # 主机内存
        host_mem = cuda.pagelocked_empty(size, dtype)
        # 设备内存
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        bindings.append(int(device_mem))
        if engine.binding_is_input(binding):
            inputs.append({"host": host_mem, "device": device_mem})
        else:
            outputs.append({"host": host_mem, "device": device_mem})
    return inputs, outputs, bindings, stream

inputs, outputs, bindings, stream = allocate_buffers(engine)

# 3. 推理函数封装
def do_inference(context, inputs, outputs, bindings, stream, data):
    # 将输入数据拷贝到主机内存
    np.copyto(inputs[0]["host"], data.ravel())
    # 主机→设备拷贝
    cuda.memcpy_htod_async(inputs[0]["device"], inputs[0]["host"], stream)
    # 执行推理
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # 设备→主机拷贝
    cuda.memcpy_dtoh_async(outputs[0]["host"], outputs[0]["device"], stream)
    stream.synchronize()  # 等待完成
    return outputs[0]["host"].reshape(1, -1)  # 调整输出形状

# 4. 测试推理
test_data = np.random.rand(1, 224, 224, 3).astype(np.float32)
pred = do_inference(context, inputs, outputs, bindings, stream, test_data)
print(f"TensorRT引擎推理结果：{np.argmax(pred)}")

三、进阶优化：INT8量化（极致加速）

FP16可提升2倍速度，INT8可提升4倍以上，但需校准数据集保证精度：

# 1. 定义INT8校准器（需准备100~1000张校准图片）
class Calibrator(trt.IInt8MinMaxCalibrator):
    def __init__(self, calib_data_path, batch_size=8, input_shape=(224,224,3)):
        trt.IInt8MinMaxCalibrator.__init__(self)
        self.batch_size = batch_size
        self.input_shape = input_shape
        self.calib_data = self.load_calib_data(calib_data_path)
        self.device_input = cuda.mem_alloc(trt.volume(input_shape) * batch_size * np.float32().nbytes)
        self.batch_idx = 0

    def load_calib_data(self, path):
        # 加载校准数据集（已归一化的numpy数组）
        calib_files = [os.path.join(path, f) for f in os.listdir(path) if f.endswith(".npy")]
        calib_data = []
        for f in calib_files[:100]:  # 取100张即可
            data = np.load(f).astype(np.float32)
            calib_data.append(data)
        return np.array(calib_data)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.batch_idx + self.batch_size > len(self.calib_data):
            return None
        batch = self.calib_data[self.batch_idx:self.batch_idx+self.batch_size]
        self.batch_idx += self.batch_size
        cuda.memcpy_htod(self.device_input, batch.ravel())
        return [int(self.device_input)]

    def read_calibration_cache(self):
        return None

    def write_calibration_cache(self, cache):
        with open("calib_cache.bin", "wb") as f:
            f.write(cache)

# 2. 构建INT8引擎
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30
config.set_flag(trt.BuilderFlag.INT8)  # 开启INT8
# 设置校准器
calibrator = Calibrator(calib_data_path="./calib_data", batch_size=8)
config.int8_calibrator = calibrator

# 后续步骤同FP16，构建引擎即可

四、性能对比（参考）

以MobileNetV2图像分类模型（输入224×224）为例，在NVIDIA T4 GPU上的性能：

推理方式精度单张推理时间吞吐量（张/秒）显存占用
TensorFlow原生FP328ms125800MB
TF-TRT（FP16）FP162ms500400MB
TensorRT引擎（INT8）INT80.8ms1250200MB
avatar
文章
阅读
粉丝