python - 对角化+可训练的自定义 Tensorflow 层

Question

我想将以下函数编码到 TS 层中。设 x 为 d 维向量。
x -> tf.linalg.diag(x)*A + b,

其中 A 是一个可训练 dxd 矩阵，b 是一个可训练(d 维)向量。

如果 A 和 b 不存在，我会使用 Lambda 层，但既然它们存在......我将如何处理它。

<小时/>

P.s.:为了教育目的，我不想提供 lambda 层:

Lambda(lambda x: tf.linalg.diag(x)))

进入具有“身份”激活的全连接层。 (我知道这有效，但它并不能帮助我学习如何真正解决问题:))

Answer 1

您可以创建自定义层并将函数放入调用方法中。

class Custom_layer(keras.layers.Layer):

    def __init__(self, dim):
        super(Custom_layer, self).__init__()
        self.dim = dim

        # add trainable weight
        self.weight = self.add_weight(shape=(dim,dim),trainable=True)
        # add trainable bias
        self.bias = self.add_weight(shape=(dim))

    def call(self, input):
        # your function
        return (tf.linalg.diag(input)*self.weight) + self.bias
    def get_config(self):
        config = super(Custom_layer, self).get_config()
        config['dim'] = self.dim
        return config

像普通图层一样使用它，并在使用时为其提供尺寸参数。

my_layer = Custom_layer(desire_dimension)
output = my_layer(input)