类型错误：无法腌制 _thread.RLock 对象 - 在 keras 中保存模型

Question

我有带有自定义层的模型。自定义图层如下

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, loaded_model):
        super(CustomLayer, self).__init__()
        self.loaded_model = loaded_model

    def build(self, input_shape):
        pass

    def get_config(self):

        config = super().get_config().copy()
        config.update({
            'loaded_model': self.loaded_model
        })
        return config

    def call(self, x):
        stacked_output = []
        for t in tf.unstack(x, axis=-1):
            single_image = tf.expand_dims(t, axis=-1)
            for i in range(10):
                single_image = self.loaded_model.get_layer(index=i)(single_image)
            stacked_output.append(single_image)
            output = tf.stack(stacked_output, axis=-1)
            # print(output.shape)
        return output

带有自定义层的主模型：

def create_model():
    trajectory_input = Input(shape=(n_steps, feature_count), name='trajectory_input')
    image_input = Input(shape=(128, 128, n_steps), name='image_input')

    x_aware = CustomLayer(loaded_model)(image_input)
    x_aware = Reshape((1, 501760))(x_aware)

    x = (Dense(32, activation='relu', kernel_regularizer=regularizers.l2(0.001)))(trajectory_input)
    x = Reshape((1, 32 * n_steps))(x)

    x = concatenate([x, x_aware])

    output_regression = (Dense(2, name='main_output'))(x_reg)

    adam = Adam(lr=learning_rate)
    model = Model(inputs=[trajectory_input, image_input], outputs=output_regression)
    model.compile(optimizer=adam, loss='mse', metrics=[euc_dist_1, euc_dist_2])
    model.summary()
    return model

当我最初尝试使用 model.save("model.h5") 保存模型时，我最终得到了

NotImplementedError: Layers with arguments in __init__ must override `get_config`

但是一旦我在 CustomLayer 中包含 get_config 就解决了错误。

但现在我明白了

TypeError: can't pickle _thread.RLock objects

我使用预训练模型作为自定义层。