4. 使用TensorRT优化GPU推理性能

任务目标

将训练好的TensorFlow模型（以ResNet50为例）转换为TensorRT引擎，并在GPU上测量推理延迟与吞吐量，对比TensorRT与原生TensorFlow的性能差异。

环境依赖

• Ubuntu 20.04+，NVIDIA GPU（本文以T4为例），Driver ≥ 470
• TensorFlow 2.10（CPU版即可，仅用于导出模型）
• TensorRT 8.6 + TF-TRT（pip install tensorflow[and-cuda]==2.10 或使用NVIDIA官方镜像 nvcr.io/nvidia/tensorflow:23.04-tf2-py3）
• PyCUDA（用于精确计时）

实现步骤

整体流程：导出SavedModel → 使用TF-TRT转换为TensorRT引擎 → 加载引擎并预热 → 测量延迟与吞吐量。TF-TRT会在运行时将支持的算子融合为TensorRT引擎，不支持的算子回退到TensorFlow执行。

完整代码

import time
import numpy as np
import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt

# ========== 1. 导出/加载TensorFlow模型 ==========
# 使用预训练ResNet50，若已有自己的模型，替换为 tf.saved_model.load 即可
model = tf.keras.applications.ResNet50(weights='imagenet', input_shape=(224, 224, 3))
model.save('/tmp/resnet50_savedmodel')

# ========== 2. 转换为TensorRT引擎 ==========
conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
    precision_mode='FP16',          # 使用FP16精度
    max_workspace_size_bytes=1 << 30,  # 1GB workspace
    minimum_segment_size=3,         # 最小算子段大小
    max_batch_size=8                # 最大批大小
)

converter = trt.TrtGraphConverterV2(
    input_saved_model_dir='/tmp/resnet50_savedmodel',
    conversion_params=conversion_params
)
converter.convert()

# 保存转换后的模型
trt_engine_path = '/tmp/resnet50_trt'
converter.save(output_saved_model_dir=trt_engine_path)
print('TensorRT引擎已保存至:', trt_engine_path)

# ========== 3. 加载引擎并构建推理函数 ==========
trt_model = tf.saved_model.load(trt_engine_path)
infer = trt_model.signatures['serving_default']

# 获取输入张量名称（ResNet50的输入名为'input_1'）
input_name = list(infer.structured_input_signature[1].keys())[0]
print('输入张量名:', input_name)

# ========== 4. 性能测试 ==========
batch_size = 8
input_shape = (batch_size, 224, 224, 3)
dummy_input = tf.random.uniform(input_shape, minval=0, maxval=1.0, dtype=tf.float32)

# 预热（Warm-up）：让TensorRT完成内核选择与显存分配
for _ in range(10):
    infer(dummy_input)

# 延迟测试：单次推理耗时（毫秒）
latencies = []
n_iter = 100
for _ in range(n_iter):
    start = time.perf_counter()
    infer(dummy_input)
    latencies.append((time.perf_counter() - start) * 1000)

latencies.sort()
avg_latency = np.mean(latencies)
p50 = latencies[n_iter // 2]
p95 = latencies[int(n_iter * 0.95)]
print(f'TensorRT 延迟 (batch={batch_size}):')
print(f'  平均: {avg_latency:.2f} ms')
print(f'  P50:  {p50:.2f} ms')
print(f'  P95:  {p95:.2f} ms')

# 吞吐量测试：每秒处理图像数
total_images = 0
start = time.perf_counter()
duration = 3.0  # 测试3秒
while time.perf_counter() - start < duration:
    infer(dummy_input)
    total_images += batch_size
elapsed = time.perf_counter() - start
throughput = total_images / elapsed
print(f'TensorRT 吞吐量: {throughput:.1f} images/sec')

# ========== 5. 对比原生TensorFlow ==========
# 加载原始模型（使用GPU）
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=4096)]
    )

tf_model = tf.keras.models.load_model('/tmp/resnet50_savedmodel')
tf_model.compile()  # 仅推理不需要compile

# 预热
for _ in range(5):
    tf_model(dummy_input, training=False)

# 延迟测试
latencies_tf = []
for _ in range(50):
    start = time.perf_counter()
    tf_model(dummy_input, training=False)
    latencies_tf.append((time.perf_counter() - start) * 1000)

avg_tf = np.mean(latencies_tf)
print(f'\n原生TensorFlow 延迟 (batch={batch_size}):')
print(f'  平均: {avg_tf:.2f} ms')

# 吞吐量测试
total_images = 0
start = time.perf_counter()
while time.perf_counter() - start < duration:
    tf_model(dummy_input, training=False)
    total_images += batch_size
elapsed = time.perf_counter() - start
throughput_tf = total_images / elapsed
print(f'原生TensorFlow 吞吐量: {throughput_tf:.1f} images/sec')

# ========== 6. 结果汇总 ==========
print('\n========== 性能对比 ==========')
print(f'延迟加速比: {avg_tf / avg_latency:.2f}x')
print(f'吞吐量提升: {throughput / throughput_tf:.2f}x')

运行方式

# 使用NVIDIA官方TensorFlow容器（已包含TF-TRT）
docker run --gpus all -it --rm \
  -v $(pwd):/workspace \
  nvcr.io/nvidia/tensorflow:23.04-tf2-py3 \
  python /workspace/trt_inference.py

验证方法与判定标准

• 转换成功标志：控制台输出“TensorRT引擎已保存至: /tmp/resnet50_trt”，且该目录下存在 saved_model.pb 与 variables/ 文件夹。
• 推理正确性：对同一输入，TensorRT与原生TF的预测类别应一致（可在代码中增加 np.argmax 对比，Top-1类别相同即通过）。
• 性能提升判定：在T4 GPU上，FP16精度下TensorRT延迟通常比原生TF快 1.5～3倍，吞吐量提升 2～4倍。若加速比低于1.2倍，检查是否使用了FP16、batch size是否≥4，以及GPU是否被其他进程占用。
• 无回退警告：转换时若出现大量“segment fallback”警告，说明模型中有较多不支持的算子，可尝试使用FP32精度或更新TensorRT版本。

常见问题

• 显存不足：减小 max_workspace_size_bytes 或 batch_size。
• 转换后精度下降：FP16可能导致个别层精度损失，若验证集准确率下降超过1%，改用FP32（precision_mode='FP32'）。
• 动态输入形状：本案例使用固定shape（224×224×3），若需动态batch，需在转换参数中设置 dynamic_shape=True 并配置优化profile。