ValueError: Layer count mismatch when loading weights from file. Model expected 39 layers, found 29 saved layers.

[CristianoLimaOliveira]

i needed to modify the srgan.py file by adding this class:

class StackLayer(tf.keras.layers.Layer):
    def call(self, inputs):
        return tf.stack(inputs)

i also needed to modify the following lines:

    if return_features:
        print(type(x_in), type(features)) # i added it for test
        # exit(0) # i added it for test
        stacked_features = StackLayer()(features) # i added here
        return tf.keras.models.Model(x_in, [x, stacked_features]) # i added here
        #return tf.keras.models.Model(x_in, [x, tf.stack(features)]) # i commented here
    else:
        return tf.keras.models.Model(x_in, x)

These changes were necessary because the following error occurred without them:

Traceback (most recent call last):
  File "/content/drive/MyDrive/EdgeSRGAN/main.py", line 45, in <module>
    trainer = Trainer(config=config, logger=logger, teacher=None, trial=None)
  File "/content/drive/MyDrive/EdgeSRGAN/utils/train.py", line 52, in __init__
    self.get_models()
  File "/content/drive/MyDrive/EdgeSRGAN/utils/train.py", line 85, in get_models
    self.generator = generator(scale=self.config['SCALE'],
  File "/content/drive/MyDrive/EdgeSRGAN/utils/srgan.py", line 114, in generator
    return tf.keras.models.Model(x_in, [x, tf.stack(features)])
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.10/dist-packages/keras/src/backend/common/keras_tensor.py", line 91, in __tf_tensor__
    raise ValueError(
ValueError: A KerasTensor cannot be used as input to a TensorFlow function. A KerasTensor is a symbolic placeholder for a shape and dtype, used when constructing Keras Functional models or Keras Functions. You can only use it as input to a Keras layer or a Keras operation (from the namespaces `keras.layers` and `keras.operations`). You are likely doing something like:

'''
x = Input(...)
'''
tf_fn(x)  # Invalid.
'''

What you should do instead is wrap `tf_fn` in a layer:

'''
class MyLayer(Layer):
    def call(self, x):
        return tf_fn(x)

x = MyLayer()(x)
'''

After these changes, when the pretrained weights from pre_generator_small.h5 are loaded into the generator, the following error is thrown:

Traceback (most recent call last):
  File "/content/drive/MyDrive/EdgeSRGAN/main.py", line 46, in <module>
    res = trainer.train()
  File "/content/drive/MyDrive/EdgeSRGAN/utils/train.py", line 211, in train
    self.generator.load_weights(pre_train_weights)
  File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.10/dist-packages/keras/src/legacy/saving/legacy_h5_format.py", line 357, in load_weights_from_hdf5_group
    raise ValueError(
ValueError: Layer count mismatch when loading weights from file. Model expected 39 layers, found 29 saved layers.