How do I adjust contrast within TensorFlow layer?

Viewed 408

I want to create a TensorFlow/Keras layer that always applies contrast, both during training time and testing time.

I used the TensorFlow contrast adjustment method tf.image.adjust_contrast, but I get an error when I place it in the call method of the layer I'm making and attempt to train:

LookupError: gradient registry has no entry for: AdjustContrastv2

This is my current code for the layer:

class Contrast(keras.layers.Layer):
    def __init__(self, contrast_level=2, **kwargs):
        super(Contrast, self).__init__(**kwargs)
        self.supports_masking = True
        self.contrast_level = contrast_level

    def call(self, inputs, training=None):
        return tf.image.adjust_contrast(inputs, self.contrast_level)

    def get_config(self):
        config = {'stddev': self.stddev}
        base_config = super(Contrast, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def compute_output_shape(self, input_shape):
        return input_shape

I figure that for efficiency I should try to use a TensorFlow method that implements contrast, but is this not the way to go about it?

1 Answers

The error message means literally, tensorflow don't know the gradient of the ops. You will have to define adjust_contrast by yourself.

Example Codes:

class Contrast(tf.keras.layers.Layer):
    def __init__(self, contrast_level=2.0, **kwargs):
        self.contrast_level = contrast_level
        super().__init__(**kwargs)
        self.supports_masking = True

    def call(self, inputs, training=None):
      channel_mean = tf.math.reduce_mean(inputs, [1, 2], keepdims=True)
      return (inputs - channel_mean) * self.contrast_level + channel_mean

contrast_layer=Contrast()

random_images=[np.random.uniform(1,2,size=(64,224,224,3)).astype(np.float32) for i in range(3)] 

print(np.allclose(tf.image.adjust_contrast(random_images[0],2),contrast_layer(random_images[0]),atol=0,rtol=1e-4))#True
print(np.allclose(tf.image.adjust_contrast(random_images[1],2),contrast_layer(random_images[1]),atol=0,rtol=1e-4))#True
print(np.allclose(tf.image.adjust_contrast(random_images[2],2),contrast_layer(random_images[2]),atol=0,rtol=1e-4))#True
Related