用于更改内核权重的包装层

Question

我正在尝试编写一个自定义包装层，例如以下（简化的），我想在其中修改包装层的内核权重：

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()

但是，代码在调用函数的第一行中断，输出如下：

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

我已经检查过https://www.tensorflow.org/addons/api_docs/python/tfa/layers/WeightNormalization，但不确定它是否有帮助。虽然他们似乎也重新定义了内核，但他们基于单独的变量而不是内核本身（在我的理解中）重新定义了它。任何帮助将不胜感激！

Answer 1

层的内核是tf.Variable。要更改其值，请使用assign 方法。

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

用张量覆盖tf.Variable 是一个常见的错误。您可以在指南中阅读有关Variabless 的更多信息：Introduction to Variables。

Variable 甚至还有一些方便的方法，比如assign_add，可以将上面的代码缩短为 self.layer.kernel.assign_add(tf.ones_like(self.layer.kernel))