TensorFlow中depth_to_space后Conv2D形状不兼容报错问题排查

你遇到的报错本质是convOut层权重初始化时的输入通道数和实际推理时的输入通道数不匹配，核心原因有2个：

• 1. Lambda层的位置错误
     你把depth_to_space对应的Lambda层写在了call方法内部，子类化Model在build阶段（权重初始化阶段）计算各层输入形状时，不会识别call内部临时定义的层的形状变换逻辑，只会统计__init__中定义的层的输出形状：convUp层的输出通道为64，因此convOut层会按照输入通道为64初始化卷积核，得到形状为(3,3,64,3)的权重，和你硬编码的input_shape要求的输入通道4冲突。
• 2. 上采样逻辑冗余
     你同时使用了2个步长为2的转置卷积（已实现4倍上采样）+ 缩放因子为4的depth_to_space，总上采样倍率会达到16倍，不符合常规4倍超分的设计逻辑，也会导致通道数变换不符合预期。

按照以下步骤修改即可解决报错：

1. 把Lambda层移到__init__中定义为模型的属性，确保build阶段可以正确识别形状变换
2. 删掉convOut层中硬编码的input_shape参数，子类化模型会自动推理输入形状，不需要手动指定
3. 修正上采样逻辑，二选一即可：
   • 保留转置卷积上采样：删掉depth_to_space相关代码，把convUp最后一层转置卷积的输出通道改为3，直接输出结果，不需要后续的convOut
   • 保留Pixel Shuffle上采样：把convUp层的输出通道改为3*scale*scale=48，删掉两个转置卷积的步长设置（改为步长1），仅用depth_to_space实现4倍上采样

class SRNet(Model):
    def __init__(self, scale=4):
        super(SRNet, self).__init__()
        self.scale = scale

        self.conv1 = Sequential([
            layers.Conv2D(filters=64, kernel_size=3,
                          strides=(1, 1), padding="same", data_format="channels_first"),
            layers.ReLU(),
        ])

        self.residualBlocks = Sequential(
            [ResidualBlock() for _ in range(16)])
        
        # 仅保留特征提取，不做转置卷积上采样，改为用Pixel Shuffle上采样
        self.convFeat = Sequential([
            layers.Conv2D(filters=64, kernel_size=3, strides=(1, 1), padding="same", data_format="channels_first"),
            layers.ReLU(),
            # Pixel Shuffle要求输入通道为 输出通道*scale^2 = 3*4*4=48
            layers.Conv2D(filters=48, kernel_size=3, strides=(1, 1), padding="same", data_format="channels_first"),
            layers.ReLU(),
        ])
        
        # Lambda层移到__init__中定义
        self.subpixel_layer = Lambda(lambda x: tf.nn.depth_to_space(
            x, self.scale, data_format="NCHW"))
        self.reluAfterPixleShuffle = layers.ReLU()

        # 删掉硬编码的input_shape
        self.convOut = layers.Conv2D(
            filters=3, kernel_size=3, strides=(1, 1), padding="same", data_format="channels_first")

    def call(self, lrCur_hrPrevTran):
        lrCur, hrPrevTran = lrCur_hrPrevTran
        x = tf.concat([lrCur, hrPrevTran], axis=1)
        x = self.conv1(x)
        x = self.residualBlocks(x)
        x = self.convFeat(x)

        x = self.subpixel_layer(inputs=x)
        x = self.reluAfterPixleShuffle(x)
        
        x = self.convOut(x)
        return x

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