Wrapper layer to change kernel weights

Question

I'm trying to write a custom wrapper layer such as the following one (simplified), where I want to modify the kernel weights of the wrapped layer:

import tensorflow as tf


class MyWrapper(tf.keras.layers.Wrapper):
    def __init__(self, layer: tf.keras.layers, **kwargs):
        super().__init__(layer, **kwargs)

    def call(self, inputs, **kwargs):
        self.layer.kernel = self.layer.kernel + 1
        outputs = self.layer(inputs)
        return outputs


def main():
    # setup model
    input_shape = (8, 8, 1)
    xin = tf.keras.layers.Input(shape=input_shape)
    xout = MyWrapper(tf.keras.layers.Conv2D(4, (3, 3), padding="same"))(xin)
    model = tf.keras.models.Model(inputs=xin, outputs=xout)
    model.compile()

    # run with output
    x_shape = (1, *input_shape)
    x = tf.random.uniform(x_shape, dtype=tf.float32)
    xout = model(x)
    print(xout)


if __name__ == "__main__":
    main()

However, the code breaks in the first line of the call function with the following output:

TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: my_wrapper/add:0

I've already checked https://www.tensorflow.org/addons/api_docs/python/tfa/layers/WeightNormalization but am not sure whether it helps. While they also seem to redefine the kernel, they redefine it based on separate variables and not the kernel itself (in my understanding). Any help would be greatly appreciated!

user_7370153 · Answer

A kernel of a layer is a tf.Variable. To change its value, use the assign method.

def call(self, inputs, **kwargs):
    self.layer.kernel.assign(self.layer.kernel + 1)
    outputs = self.layer(inputs)
    return outputs

Its a common mistake to overwrite a tf.Variable with a Tensor. You can read more about Variabless in the guide: Introduction to Variables.

A Variable even has some convenience methods like assign_add, that can shorten the code above to self.layer.kernel.assign_add(tf.ones_like(self.layer.kernel))

Wrapper layer to change kernel weights

1 Answers