DeepLab TensorFlow转TensorRT后推理耗时剧增，求正确转换方案

先说说你现有代码里可能拖慢速度的几个关键点，再给你修正后的完整流程：

一、现有代码的核心问题 #

• INT8精度未做校准：你直接用了precision_mode="INT8"，但TensorRT的INT8优化必须依赖校准数据集生成校准表，没有校准的话，模型会自动 fallback 到FP32甚至更差的性能，反而可能因为额外的转换开销变慢。
• 推理Session配置未针对TensorRT优化：你的Session只设置了per_process_gpu_memory_fraction，没有启用TensorRT相关的运行时优化，也没开GPU内存自动增长，可能导致内存分配效率低。
• 转换时未明确输入形状：TensorRT需要明确的输入张量形状才能做更充分的优化，你没指定input_shape参数，模型可能以动态形状处理，无法最大化优化。

二、正确的转换与推理流程 #

步骤1：先尝试FP16精度（快速验证提速效果） #

FP16不需要校准，只要你的GPU支持（比如Turing架构及以上），就能直接获得不错的提速，先用来验证转换是否有效：

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

OUTPUT_NAME = ["SemanticPredictions"]
INPUT_SHAPE = [2, 513, 513, 3]  # 对应max_batch_size=2，输入尺寸和你的DeepLab一致（比如513x513）

# 读取冻结图
with tf.gfile.FastGFile('/frozen_inference_graph.pb', 'rb') as tf_model:
    tf_graphf = tf.GraphDef()
    tf_graphf.ParseFromString(tf_model.read())

# 转换为TensorRT FP16模型
trt_graph = trt.create_inference_graph(
    input_graph_def=tf_graphf,
    outputs=OUTPUT_NAME,
    max_batch_size=2,
    max_workspace_size_bytes=2 * 10**9,
    precision_mode="FP16",
    input_shape=INPUT_SHAPE  # 明确输入形状
)

# 保存模型
with tf.gfile.FastGFile("TensorRT_FP16_model.pb", 'wb') as f:
    f.write(trt_graph.SerializeToString())
print("TensorRT FP16 model saved successfully!")

步骤2：优化推理脚本配置 #

推理时要给Session配置TensorRT相关的优化选项，确保GPU资源被高效利用：

import tensorflow as tf
import cv2
import numpy as np

def label_to_color_image(seg_map):
    # 你的label转颜色函数，保持不变
    pass

# 配置Session，启用TensorRT优化和GPU内存自动增长
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # 避免预分配过多GPU内存
config.gpu_options.per_process_gpu_memory_fraction = 0.50
# 启用TensorRT运行时优化（TF1.x需要这个配置）
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1

with tf.Session(config=config) as sess:
    img_array = cv2.imread('test.png', 1)
    # 确保输入尺寸和转换时一致，比如513x513
    img_array = cv2.resize(img_array, (513, 513))
    
    # 加载TensorRT模型
    with tf.gfile.FastGFile('TensorRT_FP16_model.pb', 'rb') as trt_model:
        trt_graph = tf.GraphDef()
        trt_graph.ParseFromString(trt_model.read())
    
    tf.import_graph_def(trt_graph, name='')
    input_tensor = sess.graph.get_tensor_by_name('ImageTensor:0')
    output_tensor = sess.graph.get_tensor_by_name('SemanticPredictions:0')
    
    # 先warm-up几次，消除首次推理的开销
    for _ in range(5):
        sess.run(output_tensor, feed_dict={input_tensor: [img_array]})
    
    # 正式推理
    batch_seg_map = sess.run(output_tensor, feed_dict={input_tensor: [img_array]})
    seg_map = batch_seg_map[0]
    seg_img = label_to_color_image(seg_map).astype(np.uint8)

步骤3：如果需要INT8精度（更高提速） #

要使用INT8，必须准备校准数据集（比如100-200张和测试集类似的图片），生成校准表后再转换：

# 先定义校准数据生成器
def calibration_data_generator():
    batch_size = 2
    # 替换成你的校准数据集路径列表
    calibration_images = ["calib_img1.png", "calib_img2.png", ...]
    for i in range(0, len(calibration_images), batch_size):
        batch = []
        for img_path in calibration_images[i:i+batch_size]:
            img = cv2.imread(img_path, 1)
            img = cv2.resize(img, (513, 513))
            batch.append(img)
        yield np.array(batch)

# 转换为INT8模型，传入校准生成器
trt_graph = trt.create_inference_graph(
    input_graph_def=tf_graphf,
    outputs=OUTPUT_NAME,
    max_batch_size=2,
    max_workspace_size_bytes=2 * 10**9,
    precision_mode="INT8",
    input_shape=INPUT_SHAPE,
    calibration_data=calibration_data_generator()  # 校准数据
)

# 保存INT8模型
with tf.gfile.FastGFile("TensorRT_INT8_model.pb", 'wb') as f:
    f.write(trt_graph.SerializeToString())

三、额外注意事项 #

• 版本兼容性：确保TensorFlow和TensorRT版本匹配（比如TF1.15对应TRT6.x，TF2.x对应TRT7+），版本不兼容会导致优化失效。
• 输入尺寸一致性：转换时的input_shape必须和推理时的输入尺寸完全一致，否则TensorRT的优化会失效。
• warm-up推理：首次推理会有模型加载和初始化开销，所以先跑几次warm-up再统计时间。

内容的提问来源于stack exchange，提问作者Pedram