子类化keras.Model实现Seq2Seq摘要模型训练出现function_optimizer等错误及OOM问题

问题解答 #

1. 前向代码执行两次的原因 #

这是TensorFlow自动图构建机制的正常行为：自定义训练循环首次执行时，TensorFlow会先做一次计算图追踪（trace），用来推导所有输入张量的形状、类型，生成静态计算图，第二次才会执行实际的前向计算。你使用的是Python原生print语句，两次执行过程都会触发打印，所以会看到两份完全一致的执行日志，和return语句无关。

额外的触发因素：你当前Summary_Generator的call方法实现存在参数错误：方法定义接收inputs参数，但代码中直接使用了source_seq/initial_states/dec_in等未从inputs中解析的变量，导致TensorFlow无法完成图追踪，会重复触发多次trace尝试，加剧了重复执行的问题。 #

2. 前向计算完成后仍触发OOM的原因 #

你的训练逻辑存在冗余计算+梯度带显存占用问题：

• 你在序列长度的循环内每次都完整执行了Encoder+Decoder的前向计算，Encoder每次的输出、中间张量都会被tf.GradientTape捕获留存，用于后续梯度计算，不会在前向结束后立即释放
• 真实数据集的序列长度、词表大小（你的logits维度对应词表大小为6764）远大于toy样本，多步循环累计的中间张量很快就会占满GPU显存，等到需要申请梯度计算所需的新张量时就会触发OOM报错

3. function optimizer执行失败的原因 #

报错的是TensorFlow的Grappler图优化器，触发原因就是前面提到的call方法参数实现错误：你调用模型时传入的是[text, dec_in, in_states[0], in_states[1]]列表，但call方法没有对输入列表做拆包解析，直接引用了外部变量，导致生成的计算图节点的输入输出类型不匹配，Grappler在做图优化时检测到类型冲突（报错信息中期望variant类型，实际传入float就是类型不匹配的直接表现），因此抛出优化失败的错误。 #

修复建议 #

• 修正Summary_Generator的call方法，正确解析输入参数：

def call(self, inputs):
    # 按传入顺序解析输入
    source_seq, dec_in, initial_h, initial_c = inputs
    en_outputs = self.encoder(source_seq, [initial_h, initial_c])
    en_states = en_outputs[1:]
    de_states = en_states
    print("enc call over, dec call begins")
    dec_outputs = self.decoder(dec_in, de_states, en_outputs[0])
    print("dec execution over ...") 
    return dec_outputs

• 把Encoder的计算提到循环外部，不需要每个时间步都重新跑Encoder，大幅减少冗余计算和显存占用：

with tf.GradientTape() as tape:
    # 先只跑一次Encoder拿到输出和状态
    en_outputs = sg.encoder(text, in_states)
    en_states = en_outputs[1:]
    en_feat = en_outputs[0]
    predictions = []
    for i in range((summary.shape[1])-1): 
        dec_in = tf.expand_dims(summary[:, i], 1)
        # 循环内只跑Decoder
        de_outputs = sg.decoder(dec_in, en_states, en_feat)
        predictions.append(de_outputs)

• 给train_one_step添加@tf.function装饰器，固定输入形状避免重复trace，也可将Python原生print替换为tf.print，只在实际执行时打印一次
• 适当减小batch size、截断过长的输入输出序列，也可开启混合精度训练降低显存占用

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