从TF1.0迁移至TF2 Keras：正确实现Seq2Seq的sampled_softmax_loss

我来帮你搞定这个采样softmax损失的实现问题，先梳理下你现有代码的几个核心问题，再一步步给出修正方案：

现有代码的问题分析 #

1. 损失函数依赖外部变量：你写的SampledSoftmaxLoss直接引用了外部的d_lstm_out，这不符合Keras损失函数的设计规范——损失函数应该仅依赖传入的y_true和y_pred，否则会导致计算图依赖混乱，模型保存/加载也会出问题。
2. 权重维度不匹配：原TF1代码中output_projection_w_t的维度是[vocabulary_size, state_size]，但Keras的Dense层权重默认是[input_dim, output_dim]（也就是[state_size, vocabulary_size]），直接用的话维度会对不上，需要转置。
3. 解码器结构未对齐原代码：你当前的Keras解码器只用到了单层LSTM，没有实现原代码中的多LSTM层堆叠和输出dropout逻辑。

完整修正方案 #

第一步：对齐原TF1的解码器结构 #

先重构解码器，还原原代码的多LSTM层、输出dropout逻辑：

import tensorflow as tf
from tensorflow.keras.layers import Input, Embedding, RNN, LSTMCell, DropoutWrapper
from tensorflow.keras.models import Model

# 假设你已经定义了这些参数：
vocabulary_size = 你的词汇表大小
state_size = 你的状态维度
num_lstm_layers = 原代码中的LSTM层数
tf_keep_probability = dropout保留概率
encoder_states = 编码器输出的状态（需与解码器LSTM状态维度匹配）

# 解码器输入
decoder_inputs = tf.keras.Input(shape=(None,), name='decoder_input')
# 嵌入层
emb_layer = tf.keras.layers.Embedding(vocabulary_size, state_size)
x_d = emb_layer(decoder_inputs)

# 构建带输出dropout的多LSTM层
decoder_cells = []
for _ in range(num_lstm_layers):
    lstm_cell = LSTMCell(state_size)
    # 对应原代码的DtypeDropoutWrapper（输出dropout）
    dropout_cell = DropoutWrapper(lstm_cell, output_keep_prob=tf_keep_probability)
    decoder_cells.append(dropout_cell)
# 堆叠多LSTM单元
decoder_stacked_cell = tf.keras.layers.StackedRNNCells(decoder_cells)

# 封装为RNN层，返回序列输出
decoder_rnn = RNN(decoder_stacked_cell, return_sequences=True)
d_lstm_out = decoder_rnn(x_d, initial_state=encoder_states)

# 显式定义输出投影层（对应原代码的output_projection）
# 注意：我们不会直接用这个层的输出，而是用它的权重来计算采样softmax
projection_layer = tf.keras.layers.Dense(vocabulary_size, use_bias=True)

第二步：正确实现Sampled Softmax损失函数 #

用闭包方式实现符合Keras规范的损失函数，避免依赖外部变量，同时处理权重维度：

def get_sampled_softmax_loss(projection_layer, num_sampled=500):
    def loss_fn(y_true, y_pred):
        # 获取投影层的权重并转置，匹配原TF1代码中output_projection_w_t的维度
        weights = tf.transpose(projection_layer.kernel)
        biases = projection_layer.bias
        
        # 重塑输入和标签，适配sampled_softmax_loss的要求
        flat_inputs = tf.reshape(y_pred, [-1, state_size])
        flat_labels = tf.reshape(y_true, [-1, 1])
        
        # 计算采样softmax损失
        sampled_loss = tf.nn.sampled_softmax_loss(
            weights=weights,
            biases=biases,
            labels=flat_labels,
            inputs=flat_inputs,
            num_sampled=num_sampled,
            num_classes=vocabulary_size,
            num_true=1
        )
        # 返回平均损失
        return tf.reduce_mean(sampled_loss)
    return loss_fn

第三步：构建并编译模型 #

将解码器的LSTM输出作为模型输出（因为我们要用这个输出计算采样softmax），然后使用自定义损失编译：

# 构建训练模型
training_model = Model(inputs=decoder_inputs, outputs=d_lstm_out)

# 编译模型，使用自定义采样softmax损失
training_model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=你的学习率),  # 替换为原代码的优化器
    loss=get_sampled_softmax_loss(projection_layer)
)

额外注意事项 #

1. 推理阶段的处理：原代码中区分了训练和推理解码器，Keras中你可以单独构建推理模型，使用tf.keras.layers.Lambda或者自定义循环来实现推理时的采样/贪婪解码逻辑。
2. 状态传递：确保编码器输出的encoder_states维度与解码器LSTM单元的状态维度完全匹配，如果编码器是双向LSTM，需要合并或选择其中一侧的状态传入解码器。

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