keras添加可训练的变量

因为我的网络涉及到注意力机制，实现中需要添加几个可训练的系数与各层中的特征图乘在一起，最后再求和，这样在训练后可以得到比较好各层的特征图权值。

因为用的tf2.2，感觉tf.variable不太合适，也没测试（哪位兄弟试过也可以告知一下）搜了一下发现用keras搭一个继承Layer层即可，核心很简单，但网上我找到的都是创建一个自适应的权重矩阵，而我需要的只是一个权重系数，也就是一个标量，试了很久把网上各位大佬的也没法改好，又查了下矩阵与标量的相乘函数，发现了一个问题还没人提到，有必要记录一下。

class myLinear(tf.keras.layers.Layer):
    def __init__(self):
        super(myLinear, self).__init__()


    def build(self, input_shape):
        这里不指定shape就得到一个标量，
        self.w = self.add_weight(
            initializer="random_normal",
            trainable=True,
        )

    def call(self, inputs):
        核心就是这一句，tf.scalar_mul能得到矩阵与标量乘积，不然用其他函数例如tf.matmul总会有些奇怪的报错。
        return tf.scalar_mul(self.w, inputs)

下面是官方的API，也是创建一个权重矩阵，而且把shape改成[1,1]还是不行，我才只能用tf.scalar_mul。

class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super(Linear, self).__init__()
        self.w = self.add_weight(
            shape=(input_dim, units), initializer="random_normal", trainable=True
        )
        self.b = self.add_weight(shape=(units,), initializer="zeros", trainable=True)

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b


x = tf.ones((2, 2))
linear_layer = Linear(4, 2)
y = linear_layer(x)
print(y)

tf.Tensor(
[[-0.09724902  0.04435382  0.06548684  0.1264643 ]
 [-0.09724902  0.04435382  0.06548684  0.1264643 ]], shape=(2, 4), dtype=float32)

如果有哪位兄弟知道使用官方的API创建一个标量系数，麻烦评论区告知一下，毕竟可靠点。