自定义语义分割模型转TFLite推理CONV_2D张量分配错误求助

首先明确报错核心：该错误指向编号为13的2D卷积层输入张量的通道维度（NHWC格式下最后一维）与卷积核要求的输入通道数不匹配，训练阶段动态shape可正常适配，但TFLite转换阶段的静态shape推导出现偏差。

原始报错信息：
RuntimeError: tensorflow/lite/kernels/conv.cc:349 input->dims->data[3] != filter->dims->data[3] (64 != 1)Node number 13 (CONV_2D) failed to prepare

可按以下优先级尝试修复：

• 显式固定输入shape后再转换
  自定义子类化模型如果没有提前固定输入shape，TFLite转换时会出现shape推导错误。转换前先喂入和推理尺寸一致的测试输入触发模型build，再执行转换：

  # 替换为你实际的推理输入尺寸，格式为[batch_size, height, width, channels]
  test_input = tf.random.normal([1, 512, 1024, 3])
  _ = model(test_input)
  # 保存为SavedModel后再转TFLite
  tf.saved_model.save(model, "./saved_model")
  converter = tf.lite.TFLiteConverter.from_saved_model("./saved_model")
  # 显式指定支持的算子集，避免自定义算子被错误优化
  converter.target_spec.supported_ops = [
      tf.lite.OpsSet.TFLITE_BUILTINS,
      tf.lite.OpsSet.SELECT_TF_OPS
  ]
  tflite_model = converter.convert()
  

• 排查卷积层参数配置
  优先检查编号13的CONV_2D层的groups参数，如果你自定义了分组卷积且groups值等于输入通道数（即深度可分离卷积的深度卷积部分），TFLite转换时可能误将卷积核输入通道数识别为1，和实际64通道的输入不匹配。可将该层替换为TensorFlow官方实现的tf.keras.layers.DepthwiseConv2D，规避自定义分组卷积的兼容问题。
• 排查call方法的条件分支
  如果你的子类化模型call方法中存在依赖输入尺寸、batch size的条件判断逻辑，TFLite转换时会默认推导其中一条分支的shape，和实际推理时的输入shape不匹配。可以删除非必要的条件分支，或者给转换逻辑添加输入shape签名，强制固定各层输入输出维度。
• 确认推理输入和模型输入的一致性
  检查推理阶段喂入TFLite解释器的张量shape是否和模型定义的输入shape完全匹配，尤其注意通道维度是否和训练时的输入通道数一致，避免外部预处理时的通道变换没有同步到模型计算图中。

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