使用load_model()加载包含自定义函数的模型时找不到自定义函数的解决办法

2023-10-23 09:45:48

解决加载包含自定义函数的 Keras 模型时的困难

在机器学习模型开发中，使用自定义函数扩展模型功能非常普遍。然而，在使用 load_model() 加载包含自定义函数的模型时，可能会遇到错误消息，指出找不到这些函数。这个问题很常见，了解其实用解决方案至关重要，本文将深入探讨这个问题，并提供分步指南和切实可行的解决方案。

问题根源

当我们使用 load_model() 加载包含自定义函数的模型时，错误消息可能会出现，表明找不到这些函数。这通常是由以下原因造成的：

自定义函数未序列化： 保存模型时，自定义函数不会自动序列化。
模块导入问题： 加载模型时，它无法访问定义自定义函数的模块。

解决方案

为了解决这个问题，我们可以采取以下步骤：

1. 自定义函数序列化：

在保存模型之前，我们需要将自定义函数注册为可序列化的对象。我们可以使用 CustomObjectScope 上下文管理器来实现这一点：

from tensorflow.keras.utils import CustomObjectScope

with CustomObjectScope({'custom_function': custom_function}):
    model.save('my_model.h5')

2. 模块导入：

加载模型时，我们必须确保导入定义自定义函数的模块。我们可以使用以下代码来实现：

from my_module import custom_function

model = keras.models.load_model('my_model.h5', custom_objects={'custom_function': custom_function})

示例代码：

以下代码示例演示了如何加载包含自定义函数的模型：

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.utils import CustomObjectScope

def custom_function(x):
    return tf.math.sin(x)

model = tf.keras.Sequential([
    layers.Dense(1, activation='linear'),
    layers.Lambda(custom_function)
])

with CustomObjectScope({'custom_function': custom_function}):
    model.save('my_model.h5')

from my_module import custom_function

model = tf.keras.models.load_model('my_model.h5', custom_objects={'custom_function': custom_function})

# 使用已加载的模型进行预测
predictions = model.predict(input_data)