Google Colab训练的TensorFlow LSTM模型本地加载失败求助

问题概况 #

在Google Colab（TensorFlow 2.15）借助TPU训练LSTM神经网络，将模型保存为.h5格式后，本地TensorFlow 2.17 + Keras 3.4.1环境加载时持续报错：

ValueError: Unrecognized keyword arguments passed to LSTM: {'time_major': False}

已尝试的方案均未解决：本地降级TF2.15失败、升级Colab版本丢失TPU加速、单独加载权重仍报版本错误、手动删除配置中time_major参数引发模型签名错误。

可行解决方法 #

1. 自定义LSTM类忽略不兼容参数加载模型 #

新版本Keras中LSTM已移除time_major参数，可自定义类在加载时自动丢弃该参数：

from tensorflow.keras.layers import LSTM
from tensorflow.keras.models import load_model

class CustomLSTM(LSTM):
    def __init__(self, *args, **kwargs):
        # 移除不被支持的time_major参数
        kwargs.pop('time_major', None)
        super().__init__(*args, **kwargs)

# 加载模型时指定自定义对象
model = load_model('main_model.h5', custom_objects={'LSTM': CustomLSTM})

2. 改用SavedModel格式跨版本加载 #

Colab中训练完成后，导出为TensorFlow标准SavedModel格式而非.h5：

model.save('main_model_savedmodel')

本地环境加载该格式模型（兼容性优于.h5）：

import tensorflow as tf

# 直接加载为Keras模型
model = tf.keras.models.load_model('main_model_savedmodel')

若直接加载仍有问题，可通过SavedModel接口转换为Keras模型：

saved_model = tf.saved_model.load('main_model_savedmodel')
infer_func = saved_model.signatures['serving_default']
# 构建输入输出映射
keras_model = tf.keras.Model(inputs=infer_func.inputs, outputs=infer_func.outputs)

3. 在Colab中提前清理模型配置再保存 #

训练完成后，手动移除模型配置中的time_major参数再保存，避免本地加载报错：

import tensorflow as tf
from tensorflow.keras.models import Model

# 获取当前模型配置
config = model.get_config()
# 遍历所有层，清理LSTM层的time_major参数
for layer in config['layers']:
    if layer['class_name'] == 'LSTM':
        layer['config'].pop('time_major', None)
# 用清理后的配置重建模型并加载权重
fixed_model = Model.from_config(config)
fixed_model.set_weights(model.get_weights())
# 保存处理后的模型
fixed_model.save('main_model_fixed.h5')

4. 用conda创建隔离环境安装TF2.15 #

本地pip安装TF2.15失败通常是Python版本不兼容（TF2.15仅支持Python3.8~3.10），用conda创建隔离环境解决：

# 创建名为tf215的隔离环境，指定Python3.10
conda create -n tf215_env python=3.10
# 激活环境
conda activate tf215_env
# 安装TF2.15
pip install tensorflow==2.15

在该环境中加载原.h5模型即可正常运行。

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