Tensorflow - incompatible shape with conv2d after depth to space-CodePudding

I'm having a problem implementing a super-resolution model

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)])

        self.convUp = Sequential([
            layers.Conv2DTranspose(filters=64, kernel_size=3, strides=(
                2, 2), padding="same", data_format="channels_first"),
            layers.ReLU(),
            layers.Conv2DTranspose(filters=64, kernel_size=3, strides=(
                2, 2), padding="same", data_format="channels_first"),
            layers.ReLU(),
        ])

        self.reluAfterPixleShuffle = layers.ReLU()

        self.convOut = layers.Conv2D(
            filters=3, kernel_size=3, strides=(1, 1), padding="same", data_format="channels_first", input_shape=(4, 1440, 2560))  # (kernel, kernel, channel, output)

    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.convUp(x)

        # pixel shuffle
        Subpixel_layer = Lambda(lambda x: tf.nn.depth_to_space(
            x, self.scale, data_format="NCHW"))
        x = Subpixel_layer(inputs=x)
        x = self.reluAfterPixleShuffle(x)
        
        x = self.convOut(x)
        return x

Error

/usr/src/app/generator.py:164 call  *
        x = self.convOut(x)
ValueError: Tensor's shape (3, 3, 64, 3) is not compatible with supplied shape (3, 3, 4, 3)

after reading the error I know that (3, 3, 4, 3) is (kernel size, kernel size, channel, output) mean that only channel of input is not correct

so I printed out the shape of the input

# after pixel shuffle before convOut
print(x.shape)
>>> (1, 4, 1440, 2560) (batch size, channel, height, width)

but the shape of x after pixel shuffle (depth_to_space) is (1, 4, 1440, 2560) the channel value is 4 which is the same as convOut need

question is why the input's channel is changing from 4 to 64 as the error?

CodePudding user response：

I have found a solution

First of all, I'm using checkpoints to save model weight when training
during the implementation and testing of the model, I have changed some of the layers so the input size is changed too, but my weight still remember the input size from the previous checkpoint

so I delete the checkpoints folder and then everything works again