使用TensorRT进行图像超分未达预期结果求助

结合你已经跑通MNIST示例，但自定义超分网络出问题的情况，咱们从几个核心方向一步步排查：

1. 优先检查输入数据的维度与格式匹配 #

这是TRT和TensorFlow交互时最容易踩的坑！

• 你的TensorFlow模型用的是HWC格式（高度、宽度、通道），但看TRT输入注册代码：

  parser.register_input(self.input_node, (self.channels, self.height, self.width), 0)
  

  这里注册的是CHW格式（通道、高度、宽度），两者维度顺序完全相反！
• 你的输入代码里，img = img[0]后得到的是(240, 320, 3)的HWC数据，直接传给TRT的话，网络会把通道维度当成高度，完全读错输入。
• 修复建议：在传入infer前，先把数据转成CHW格式，再加上batch维度：

  img = img.transpose(2, 0, 1)  # HWC -> CHW
  img = np.expand_dims(img, axis=0)  # 增加batch维度，变成(1, 3, 240, 320)
  

2. 验证数据类型与归一化/反归一化的一致性 #

你的输出用了tanh激活，输出范围是[-1, 1]，要确保TF和TRT的前后处理完全一致：

• 检查normalize函数：是不是把0-255的uint8图像转成了[-1, 1]的float？TRT的输入是否严格遵循这个范围？
• 检查inormalize函数：TRT输出的tanh结果是不是正确转回到0-255的uint8？比如TF里如果是(x + 1) * 127.5，TRT的反归一化必须完全一样，不能有精度偏差。
• 确认TRT engine的数据类型：如果你的TF模型用的是float32，要确保TRT没有默认启用FP16（除非你特意配置且处理了精度），可以在创建engine时显式指定：

  builder = trt.Builder(g_logger)
  builder.set_fp16_mode(False)  # 先以float32调试，排除精度问题
  

3. 排查UFF模型转换的正确性 #

UFF转换TF模型到TRT时，某些层（比如反卷积）可能会被错误解析：

• 转换时开启详细日志，把LogSeverity.ERROR改成LogSeverity.INFO或LogSeverity.VERBOSE，看看有没有关于deconv层的警告或错误：

  g_logger = trt.infer.ConsoleLogger(trt.infer.LogSeverity.VERBOSE)
  

• 用UFF工具查看转换后的模型结构，确认反卷积层的参数（stride、kernel size、scale）和TF里的定义完全一致。如果是自定义的deconv2d函数，可能需要在UFF转换时注册自定义解析器。

4. 检查GPU内存拷贝与输出张量的形状处理 #

• 确认self.output_size的计算是否正确：输出应该是batch_size * 3 * 480 * 640（超分2倍后的CHW格式总元素数），如果计算错误，会导致内存拷贝不完整或越界。
• 你的输出reshape代码存在顺序问题：TRT的输出是CHW格式的一维数组，直接reshape会导致数据乱序。正确的处理步骤应该是：

  # 先把一维数组转成CHW的3D张量
  result = result.reshape(3, 480, 640)
  # 转成HWC格式再增加batch维度
  result = result.transpose(1, 2, 0)[np.newaxis, ...]
  

5. 验证TRT优化是否破坏网络 #

如果前面的步骤都没问题，可以暂时关闭TRT的优化，看看结果是否正常：

# 在创建engine时添加禁用优化的选项
builder = trt.Builder(g_logger)
builder.set_flag(trt.BuilderFlag.DISABLE_OPTIMIZATION)

如果禁用优化后结果正常，说明是TRT的某个优化步骤（比如层融合、常量折叠）出了问题，再逐步开启优化，定位具体是哪个优化导致的异常。

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