TF no gradients for variable or no gradient flow

Viewed 32

I am trying to compute gradients for the trainable parameters of convolution filter shape. The part of the code of generating the filter shape from the tf.Variable seems to work fine. But when I'm trying to compute gradients I see that they are None for corresponding tf.Variable and if I apply these None gradients I get ValueError: No gradients provided for any variable. It seems that I have no gradient flow somehow. I can not get why. Suggestions please?

Code to recreate the problem:

import tensorflow as tf

@tf.function
def mean_filter2d_p(image, filter_shape=(3, 3), name=None):
    with tf.name_scope(name or "mean_filter2d"):
        image = tf.convert_to_tensor(image, name="image")

        rank = image.shape.rank
        if rank != 3 and rank != 4:
            raise ValueError("image should be either 3 or 4-dimensional.")

        # Expand to a 4-D tensor
        if rank == 3:
            image = tf.expand_dims(image, axis=0)

        area = tf.cast(filter_shape[0] * filter_shape[1], dtype=tf.float32)
        filter_shape = tf.concat([filter_shape, [tf.shape(image)[-1], 1]], axis=-1)
        kernel = tf.ones(shape=filter_shape, dtype=image.dtype) / area

        output = tf.nn.depthwise_conv2d(image, kernel, strides=(1, 1, 1, 1), padding="SAME")

        if rank == 3:
            output = tf.squeeze(output, axis=0)

        return output

def highpass_median_p(images, filter_shape=(3, 3)):
    blurred = mean_filter2d_p(image=images, filter_shape=filter_shape)
    subtracted = images - blurred

    return subtracted, images, blurred

resolution = 128
img = tf.ones([1, resolution, resolution, 3])


filter_shapew = tf.Variable([3.0 / 128.0, 3.0 / 128.0], trainable=True, name='filter_shape')


with tf.GradientTape() as tape:

    filter_shape = tf.math.sigmoid(filter_shapew * tf.cast(resolution, dtype=tf.float32), name='filter_shape_clip')
    filter_shape = tf.math.minimum(tf.cast(filter_shape, dtype=tf.int32) + 1, resolution)

    x, x_o, x_b = highpass_median_p(img, filter_shape)

    loss_value = 1.0 - tf.reduce_mean(tf.reshape(x, [-1]))
    print(loss_value)

grads = tape.gradient(loss_value, filter_shapew)
print(grads)
0 Answers
Related