在 Keras 中，如何保存和加载包含自定义损失函数的神经网络模型？

Question

当我使用自定义损失函数时，我在保存和重新加载神经网络模型时遇到困难。 例如，在下面的代码中（集成了相关问题here和here的建议），“Save/Load Attempt 0”可以正常工作，而“Save/Load Attempt 1”则不会，返回神秘的无论模型是否加载参数

TypeError: string indices must be integers, not 'str'

或

compile=False

，都会出现错误

custom_objects={'loss': custom_loss}

。 如何修改“Save/Load Attempt 1”才能成功？

import os
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
from keras.models import Sequential, load_model
from keras.layers import Input, Dense
from keras import ops

path = 'C:/Users/.../AppData/Local/Programs/Python/Python312/.../'  # The ... represent single folders.

# ----------------------------------------------------------------------------------------------------
# Save/Load Attempt 0

dnn = Sequential()
dnn.add(Input(shape=(3,)))
dnn.add(Dense(units=5, activation='relu'))
dnn.add(Dense(units=1))
dnn.compile(loss='mean_absolute_error', optimizer='adam')

model_path_0 = path + 'dnn_0.h5'
dnn.save(model_path_0)

dnn = load_model(model_path_0)

# ----------------------------------------------------------------------------------------------------
print('---')
# Save/Load Attempt 1

def custom_loss(y_true, y_pred):
    squared_difference = ops.square(y_true - y_pred)
    return ops.mean(squared_difference, axis=-1)  # flattens squared_difference

dnn = Sequential()
dnn.add(Input(shape=(3,)))
dnn.add(Dense(units=5, activation='relu'))
dnn.add(Dense(units=1))
dnn.compile(loss=custom_loss, optimizer='adam')

model_path_1 = path + 'dnn_1.h5'
dnn.save(model_path_1)

dnn = load_model(model_path_1, compile=False)
# dnn = load_model(model_path_1, custom_objects={'loss': custom_loss})

# ----------------------------------------------------------------------------------------------------

---

供参考，错误的回溯如下。

Traceback (most recent call last):
  File "c:\Users\...\AppData\Local\Programs\Python\Python312\...\test_load_model.py", line 43, in <module>
    dnn = load_model(model_path_1, compile=False)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\...\AppData\Local\Programs\Python\Python312\Lib\site-packages\keras\src\saving\saving_api.py", line 183, in load_model
    return legacy_h5_format.load_model_from_hdf5(filepath)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\...\AppData\Local\Programs\Python\Python312\Lib\site-packages\keras\src\legacy\saving\legacy_h5_format.py", line 155, in load_model_from_hdf5
    **saving_utils.compile_args_from_training_config(
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\...\AppData\Local\Programs\Python\Python312\Lib\site-packages\keras\src\legacy\saving\saving_utils.py", line 145, in compile_args_from_training_config
    loss = _resolve_compile_arguments_compat(loss, loss_config, losses)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\...\AppData\Local\Programs\Python\Python312\Lib\site-packages\keras\src\legacy\saving\saving_utils.py", line 245, in _resolve_compile_arguments_compat
    obj = module.get(obj_config["config"]["name"])
                     ~~~~~~~~~~^^^^^^^^^^
TypeError: string indices must be integers, not 'str'

Answer 1

您是否

compiling

模型重新加载后？你没有提到你正在使用什么版本的

keras

，但是，我写了一个使用tensorflow和keras

2.14.0

以及tensorflow

2.17.0

和keras

3.5.0

测试的玩具示例：

import numpy as np
import tensorflow as tf

from tensorflow.keras import Sequential
from keras.layers import Input, Dense

# Model architecture:
model = Sequential()
model.add(Input(shape=(100,)))
model.add(Dense(units=5, activation='relu'))
model.add(Dense(units=1))

# Custom Loss
def custom_loss(y_true, y_pred):
    squared_difference = tf.math.square(y_true - y_pred)
    return tf.reduce_mean(squared_difference, axis=-1)  

# Compile model:
model.compile(optimizer="adam", loss=custom_loss, metrics= ["mean_squared_error"])

# Show summary:
model.summary()

# Some random inputs/targets:
x=np.random.rand(300,100)
y=np.random.rand(300,2)

# Fit the model for 5 epochs:
model.fit(x,y,batch_size=100, epochs=5)

# Save model
path = 'saved_model/myModel.keras'
model.save(path)
print("Model saved")

# Load model:
model = tf.keras.models.load_model(path, compile=False, custom_objects={"custom_loss": custom_loss})
print("Model reloaded")

# Compile:
model.compile(optimizer="adam", loss=custom_loss, metrics= ["mean_squared_error"])

# Continue training for 5 more epochs:
model.fit(x, y, batch_size=100, epochs=5)
print("Done fitting")

输出：

Model: "sequential_6"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 dense_12 (Dense)            (None, 5)                 505       
                                                                 
 dense_13 (Dense)            (None, 1)                 6         
                                                                 
=================================================================
Total params: 511 (2.00 KB)
Trainable params: 511 (2.00 KB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________
Epoch 1/5
3/3 [==============================] - 1s 8ms/step - loss: 0.5925 - mean_squared_error: 0.5925
Epoch 2/5
3/3 [==============================] - 0s 10ms/step - loss: 0.4522 - mean_squared_error: 0.4522
Epoch 3/5
3/3 [==============================] - 0s 8ms/step - loss: 0.3607 - mean_squared_error: 0.3607
Epoch 4/5
3/3 [==============================] - 0s 10ms/step - loss: 0.2928 - mean_squared_error: 0.2928
Epoch 5/5
3/3 [==============================] - 0s 15ms/step - loss: 0.2372 - mean_squared_error: 0.2372
Model saved
Model reloaded
Epoch 1/5
3/3 [==============================] - 1s 6ms/step - loss: 0.2142 - mean_squared_error: 0.2142
Epoch 2/5
3/3 [==============================] - 0s 4ms/step - loss: 0.1967 - mean_squared_error: 0.1967
Epoch 3/5
3/3 [==============================] - 0s 4ms/step - loss: 0.1827 - mean_squared_error: 0.1827
Epoch 4/5
3/3 [==============================] - 0s 5ms/step - loss: 0.1659 - mean_squared_error: 0.1659
Epoch 5/5
3/3 [==============================] - 0s 4ms/step - loss: 0.1528 - mean_squared_error: 0.1528
Done fitting

一些注意事项： 推荐的模型保存格式是使用

.keras

扩展名。我在这里只是假设这是一个回归问题，但这并不重要，因为数据是随机的，我们只是测试模型是否正确保存/重新加载，只需要检查指标是否确实在训练留下的地方拾取.