TensorFlow1.8升2.4后LAS模型测试加载Checkpoint参数缺失问题

• TensorFlow版本：2.4.0
• Python版本：3.6.2
• 操作背景：将LAS模型从TensorFlow 1.8.0升级至2.4.0，训练环节运行正常，测试阶段加载模型时报参数缺失，但已确认保存的模型文件中存在BeamSearchDecoderStep/multi_rnn_cell/cell_0_attention/attention_wrapper/lstm_cell_9/bias参数

1. 变量命名空间差异问题 #

TF2中tf.keras.layers.RNN、AttentionWrapper的默认变量命名规则和TF1的tf.contrib.rnn系列API存在序号后缀差异，训练时如果用兼容模式（如tf.compat.v1接口）保存变量，测试时用原生TF2 API构建模型加载，会出现命名序号不匹配的情况，比如TF1生成的lstm_cell_9后缀，在TF2原生实现里可能自动生成lstm_cell、lstm_cell_1这类连续序号，和保存的参数序号无法对应。

2. BeamSearchDecoder作用域绑定问题 #

训练阶段如果没有显式给BeamSearchDecoder添加和推理阶段一致的变量作用域，训练时保存的参数会自动带上训练阶段的临时作用域前缀，推理阶段构建的BeamSearchDecoder作用域和训练时不匹配，就算变量后缀一致也会识别为参数缺失。需要在训练和推理的BeamSearchDecoder外层统一套with tf.name_scope("BeamSearchDecoderStep"):或者with tf.variable_scope("BeamSearchDecoderStep", reuse=tf.AUTO_REUSE):，保证作用域完全一致。

3. 调整ckpt加载匹配逻辑 #

如果你用的是tf.train.Checkpoint加载参数，默认是严格匹配模式，可以先改成checkpoint.restore(ckpt_path).expect_partial()先忽略非关键缺失参数，打印出具体缺失和不匹配的参数列表，对比已保存参数的命名和当前模型生成的变量命名差异，确认是前缀、后缀还是序号的问题后，用tf.train.init_from_checkpoint手动做变量名映射，示例代码如下：

# 手动变量映射示例，根据实际命名差异调整匹配规则
var_map = {}
for var in tf.trainable_variables():
    ckpt_var_name = var.name.split(':')[0].replace('lstm_cell', 'lstm_cell_9')
    var_map[ckpt_var_name] = var
tf.train.init_from_checkpoint(ckpt_path, var_map)

4. LSTMCell实现一致性校验 #

TF1的tf.contrib.rnn.LSTMCell和TF2的tf.keras.layers.LSTMCell的变量存储结构存在差异，如果你训练时用的是compat.v1的LSTMCell，推理时用了原生TF2的LSTMCell，就算变量名一致也会识别为参数不匹配，需要保证训练和推理用的LSTMCell实现完全一致，推荐统一用tf.compat.v1.nn.rnn_cell.LSTMCell搭配兼容模式的AttentionWrapper，避免API实现差异。

内容的提问来源于stack exchange，提问作者YR123-123