Using custom, non-trainable layers in the middle of Keras model yields gradient error

[micahreich]

Problem

I am writing a GAN-style model which uses an untrained / weightless layer in the middle of the GAN (a transformation is performed on the generator output which is then fed into the discriminator). The transformation operation is done with a custom layer, but I'm getting the following error at runtime when trying to train (after compiling both the discriminator and GAN model):

ValueError: No gradients provided for any variable: ['dense_1/kernel:0', 'dense_1/bias:0', 'dense_2/kernel:0', 'dense_2/bias:0', 'dense_3/kernel:0', 'dense_3/bias:0', 'dense_4/kernel:0', 'dense_4/bias:0', 'dense_5/kernel:0', 'dense_5/bias:0'].

Code

My Custom Layer code is as follows (is meant to convert top left, bottom right coordinates into a square shape):

class ImagePaste(tf.keras.layers.Layer):
    def __init__(self,  **kwargs):
        super(ImagePaste, self).__init__(**kwargs)
        self.canvas_size = 72

    def build(self, input_shape):
        self.batch_s = input_shape[0]

        super(ImagePaste, self).build(input_shape)
    
    def call(self, input_data, t_val=255):
        positions = tf.convert_to_tensor(input_data)
        positions = tf.reshape(positions, [-1, 2, 2])

        canvas = tf.Variable(tf.ones([self.batch_s, 72, 72, 3]) * t_val,
                             dtype=tf.float32,
                             trainable=False, validate_shape=True)

        for i in tf.range(0, self.batch_s):
            color = tf.convert_to_tensor(np.eye(3)[np.random.choice(3)] * 250, dtype=tf.float32)
            tl, br = tf.cast(positions[i], dtype=tf.int64)

            for r in tf.range(min(tl[0], self.canvas_size), min(br[0], self.canvas_size)):
                for c in tf.range(min(tl[1], self.canvas_size), min(br[1], self.canvas_size)):
                    canvas[i, r, c, :].assign(color)

        return tf.convert_to_tensor(canvas)

My Generator code is as follows (turns latent space into top left, bottom right square coordinates then image of square on a canvas):

    def build_composer(self):
        latent_input = Input(shape=100, batch_size=64)

        p = Dense(units=128)(latent_input)
        p = Dense(units=256)(p)
        p = Dense(units=512)(p)
        p = Dense(units=1024)(p)

        out = Dropout(0.4)(p)
        out = Dense(units=4, activation='relu')(out)  # TOP LEFT, BOTTOM RIGHT COORDINATE

        composed_image = ImagePaste(trainable=False, dynamic=True)(out)

        return tf.keras.Model(inputs=latent_input, outputs=composed_image)

My Generator + Discriminator (GAN) code is as follows:

    def build_full_model(self, composer, discriminator):
        latent_input = Input(shape=100, batch_size=64)

        composed_image = composer(latent_input)

        discriminator.trainable = False
        valid = discriminator(composed_image)

        return tf.keras.Model(inputs=latent_input, outputs=valid)

And the error shown above occurs when calling .predict() on the compiled GAN model.

It seems like my ImagePaste Layer is preventing gradients from reaching the generator layers -- how can I get the gradient calculation to ignore the layer (even though its already set to trainable = False) when training? Can anyone help me solve this? Thank you very much.

[Saduf2019]

@micahreich
Please fill the issue template, it helps us to replicate the issue on the tf version you have faced the error on along with the system information.

Please provide complete code such that we can replicate the issue(with all dependencies) or if possible share a colab gist with error faced.

With respect to the error reported, refer to below resolved issues with same error:
#39169 #41162 #42038 link link1

[micahreich]

I created a new issue according to the correct template @Saduf2019:
Here is the new issue

Im sorry for any confusion, thank you.

[Saduf2019]

@micahreich
Moving this issue to closed status as #42311 would track the issue, also please let us know if you followed the links and issues shared.

[micahreich]

Thanks -- I looked at the links you posted but I wasn't able to make a fix based on those threads