如何使用TensorFlow正确保存并加载Encoder-Decoder编解码模型？

问题根因 #

• 尝试1失败原因：H5格式仅保存层的权重和基本结构，不会保存推理子模型和原训练模型之间的权重绑定关系，推理模型依赖的共享层权重在加载时没有正确关联，导致预测结果完全错误。
• 尝试2失败原因：你直接将Input层对象当做张量传入Functional模型构造函数，Functional模型要求输入必须是tf.keras.Input生成的张量节点，不是层对象本身。

解决方法 #

方案1：优先使用SavedModel格式保存加载（推荐）

SavedModel是TensorFlow原生格式，可完整保留计算图结构和权重绑定关系，无需重建模型即可直接使用：

1. 保存阶段修改代码：

# 去掉.h5后缀，默认使用SavedModel格式
model.save('training_model')
encoder_model.save('encoder_model')
decoder_model.save('decoder_model')

with open('tokenizer.pickle', 'wb') as handle:
    pickle.dump(tokenizer, handle, protocol=pickle.HIGHEST_PROTOCOL)

2. 加载阶段直接调用即可：

import tensorflow as tf
import pickle
import numpy as np

# 加载模型和分词器
encoder_model = tf.keras.models.load_model('encoder_model')
decoder_model = tf.keras.models.load_model('decoder_model')
with open('tokenizer.pickle', 'rb') as f:
    tokenizer = pickle.load(f)

# 直接使用原有decode_sequence函数预测即可，准确率和训练后一致

方案2：基于H5格式的训练模型重建推理模型

如果必须使用已保存的H5格式训练模型，可通过复用训练模型层权重、重新构造推理子图的方式解决：

import tensorflow as tf
import pickle
import numpy as np

# 加载训练模型和分词器
training_model = tf.keras.models.load_model('training_model.h5')
with open('tokenizer.pickle', 'rb') as f:
    tokenizer = pickle.load(f)

# 提取训练模型的共享层，复用权重
encoder_embedding = training_model.layers[2]
encoder_lstm = training_model.layers[4]
decoder_embedding = training_model.layers[3]
decoder_lstm = training_model.layers[5]
decoder_dense = training_model.layers[6]

# 重建编码器推理模型
encoder_inputs = tf.keras.Input(shape=(None,))
en_x = encoder_embedding(encoder_inputs)
encoder_outputs, state_h, state_c = encoder_lstm(en_x)
encoder_states = [state_h, state_c]
encoder_model = tf.keras.Model(encoder_inputs, encoder_states)

# 重建解码器推理模型
decoder_inputs = tf.keras.Input(shape=(None,))
decoder_state_input_h = tf.keras.Input(shape=(512,))
decoder_state_input_c = tf.keras.Input(shape=(512,))
decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]

final_dex = decoder_embedding(decoder_inputs)
decoder_outputs2, state_h2, state_c2 = decoder_lstm(final_dex, initial_state=decoder_states_inputs)
decoder_states2 = [state_h2, state_c2]
decoder_outputs2 = decoder_dense(decoder_outputs2)

decoder_model = tf.keras.Model(
    [decoder_inputs] + decoder_states_inputs,
    [decoder_outputs2] + decoder_states2
)

# 直接使用原有decode_sequence函数预测即可

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