How to modify tensor inside custom layer tensorflow?-CodePudding

I have this custom layer that I get an error when I use it inside the model.

class NormalLayer(tf.keras.layers.Layer):
  def __init__(self, color_space):
    super(NormalLayer, self).__init__()
    self.color_space = color_space

  def call(self, image):
    if self.color_space=="HSV":
        image = tf.Variable(image)
        image[:, :,0]=tf.divide(image[:, : ,0],180.)
        image[:, :,1:3]=tf.divide(image[:, :, 1:3],255.)
        return tf.convert_to_tensor(image)
    return tf.keras.layers.Rescaling(1./255)(image)

This is the error:

File "<ipython-input-55-df46ae690c0b>", line 15, in call  *
    image = tf.Variable(image)

    ValueError: Argument `initial_value` (Tensor("Placeholder:0", shape=(None, 224, 224, 3), dtype=float32)) 
    could not be lifted out of a `tf.function`. (Tried to create variable with name='{name}'). 
    To avoid this error, when constructing `tf.Variable`s inside of `tf.function` you can create the `initial_value` tensor in a `tf.init_scope` or pass a callable `initial_value` (e.g., `tf.Variable(lambda : tf.truncated_normal([10, 40]))`). 
    Please file a feature request if this restriction inconveniences you.
  
Call arguments received:
  • image=tf.Tensor(shape=(None, 224, 224, 3), dtype=float32)

Please modify this code for me.

CodePudding user response：

Try something like this:

import numpy as np
import tensorflow as tf

class NormalLayer(tf.keras.layers.Layer):
  def __init__(self, color_space):
    super(NormalLayer, self).__init__()
    self.color_space = color_space

  def call(self, image):
    if self.color_space=="HSV":
        return tf.concat([tf.divide(image[:, :, : ,0, tf.newaxis], 180.), tf.divide(image[:, :, :, 1:3],255.)], axis=-1)
    return tf.keras.layers.Rescaling(1./255)(image)

norm_layer = NormalLayer("HSV")
image = tf.random.normal((1, 244, 244, 3))
print(norm_layer(image).shape)

(1, 244, 244, 3)