在keras中重用函数的权重

import tensorflow as tf
from tensorflow.keras.layers import Layer

class SharedLayer(Layer):
    def __init__(self, units, **kwargs):
        super(SharedLayer, self).__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.kernel = self.add_weight(name='kernel', shape=(input_shape[1], self.units), initializer='uniform', trainable=True)

    def call(self, inputs):
        return tf.matmul(inputs, self.kernel)

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input

# 创建输入层
input_layer = Input(shape=(10,))

# 创建第一个模型
shared_layer = SharedLayer(units=5)
output1 = shared_layer(input_layer)
model1 = Model(inputs=input_layer, outputs=output1)

# 创建第二个模型
output2 = shared_layer(input_layer)
model2 = Model(inputs=input_layer, outputs=output2)

model1.compile(optimizer='adam', loss='mse')
model1.fit(x_train, y_train, epochs=10)

model2.compile(optimizer='adam', loss='mse')
model2.fit(x_train, y_train, epochs=10)