TensorFlow/Keras重新编译模型后GPU停止使用问题排查

你遇到的重新编译后训练速度骤降的问题，核心原因大概率是重新编译会让TensorFlow丢弃之前已经优化好的GPU计算图和kernel缓存，需要完全重新构建训练计算流程——新的计算图不仅要重新适配设备，还可能因为缺少之前的优化缓存，导致运行效率大幅下降。甚至可能出现设备放置错误，部分核心操作跑到CPU上执行，直接拖慢了整体速度。

另外，你直接修改optimizer.lr = 0.001的方式也存在隐患：在TensorFlow 2.x中，Adam优化器的学习率是一个可训练的Variable，直接赋值可能不会被计算图正确识别；而且重新编译时，优化器的内部状态（比如动量、一阶矩/二阶矩累积值）可能无法正确迁移到新计算图中，导致优化器本身跑到CPU上运行，进一步加剧速度问题。

至于直接修改model.loss_weights无效，是因为Keras的loss_weights是编译时的固化配置，一旦编译完成就会嵌入到计算图里， runtime修改不会生效，必须重新编译——这也是你当前困境的根源。

方案1：动态调整损失权重，避免重新编译（推荐） #

既然重新编译是问题的核心，我们可以通过自定义损失函数，用可训练变量动态控制损失权重，这样不需要重新编译模型就能切换训练阶段：

# 用tf.Variable存储损失权重，支持runtime动态更新（设置trainable=False避免被优化）
loss_weights = tf.Variable([0.7, 0.3], dtype=tf.float32, trainable=False)

# 自定义带权重的损失函数
def weighted_categorical_crossentropy(y_true, y_pred):
    base_loss = tf.keras.losses.categorical_crossentropy(y_true, y_pred)
    return loss_weights[0] * base_loss

def weighted_binary_crossentropy(y_true, y_pred):
    base_loss = tf.keras.losses.binary_crossentropy(y_true, y_pred)
    return loss_weights[1] * base_loss

# 编译模型时使用自定义损失
optimizer = tf.keras.optimizers.Adam(lr=0.01, clipvalue=0.1)
model.compile(
    optimizer=optimizer,
    loss=[weighted_categorical_crossentropy, weighted_binary_crossentropy],
    sample_weight_mode="temporal",
    metrics=['accuracy']
)

# 第一阶段训练
model.fit(
    train_generator(),
    steps_per_epoch=BATCHES_PER_EPOCH,
    epochs=375,
    callbacks=[reduce_lr]
)

# 切换到第二阶段：直接更新损失权重和学习率，无需重新编译
loss_weights.assign([0.99, 0.01])
optimizer.lr.assign(0.001)

# 继续第二阶段训练
model.fit(
    train_generator(),
    steps_per_epoch=BATCHES_PER_EPOCH,
    epochs=125,
    callbacks=[reduce_lr]
)

这种方式完全保留了原来的计算图和GPU优化，不会出现速度下降的问题，同时完美实现了两阶段的损失权重切换。

方案2：重新编译时确保优化器状态正确迁移 #

如果你必须重新编译模型，可以通过以下步骤避免设备放置问题：

1. 不要复用旧优化器实例，创建新优化器并迁移旧状态：

# 第一阶段训练结束后，保存旧优化器的权重（保留训练动量等状态）
old_opt_weights = optimizer.get_weights()

# 创建新的优化器，设置新的学习率
new_optimizer = tf.keras.optimizers.Adam(lr=0.001, clipvalue=0.1)
# 让新优化器初始化变量（通过build模型的可训练变量）
new_optimizer.build(model.trainable_variables)
# 加载旧优化器的状态，延续之前的训练进度
new_optimizer.set_weights(old_opt_weights)

# 强制在GPU上编译模型，避免设备放置错误
with tf.device('/GPU:0'):
    model.compile(
        optimizer=new_optimizer,
        loss=['categorical_crossentropy', 'binary_crossentropy'],
        sample_weight_mode="temporal",
        loss_weights=[0.99, 0.01],
        metrics=['accuracy']
    )

2. 替换fit_generator为model.fit：TensorFlow 2.x中fit_generator已经被弃用，fit可以直接接受生成器，兼容性更好，能减少设备放置错误的概率。

额外检查 #

你可以用nvidia-smi命令查看训练时的GPU利用率，如果第二阶段利用率明显下降，说明确实有操作跑到CPU上了。也可以用tf.debugging.assert_equal(tf.device(model.layers[0].weights[0].device), '/GPU:0')来检查模型变量是否正确放置在GPU上。

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