Keras含自定义损失函数的模型存加载失败问题及解决

在Keras中使用自定义损失函数时，内置损失函数的模型可正常保存加载，但使用自定义损失函数的模型，无论通过custom_objects参数传递损失函数，还是设置compile=False加载，都会抛出TypeError: string indices must be integers, not 'str'错误。

复现代码 #

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/.../'  # 省略具体文件夹路径

# ----------------------------------------------------------------------------------------------------
# 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)  # 对平方差取均值

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, custom_objects={'loss': custom_loss})
# dnn = load_model(model_path_1, compile=False)

错误回溯信息 #

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'

解决方法 #

通过以下三步解决该问题：

• 卸载旧版本Keras（3.1.1）并安装稳定版（3.5.0）：

  pip uninstall keras
  pip install keras
  

• 将模型存储格式从.h5改为Keras官方推荐的新格式.keras；
• 加载模型时设置compile=False，加载完成后手动执行dnn.compile()命令，重新指定自定义损失函数和优化器。

内容的提问来源于stack exchange，提问作者SapereAude