加载h5格式TensorFlow预训练模型训练报无法分配设备错误如何解决？

问题原因 #

• 本质是Apple Metal平台的TensorFlow适配限制：从报错的共址调试信息可以看到，Equal、FloorDiv、Mul三类操作在你当前的tensorflow-macos环境下仅支持CPU运行，但模型关联的Lookahead优化器节点、卷积层操作被框架默认绑定到GPU，同一计算组的操作无法同时适配两种设备，触发设备分配冲突。
• 与你提到的自定义类直接相关：你使用的h5文件不仅保存了模型权重，还保留了训练阶段绑定的自定义优化器（Lookahead）计算逻辑，即便加载时设置了compile=False，微调调用fit时，旧优化器的设备绑定规则依然会生效，与Metal GPU的适配规则冲突。
• 无原始代码的限制导致你无法直接修改自定义优化器的实现来适配Metal平台，需要通过无侵入的方式绕过冲突。

无需重新预训练的解决方案 #

方案1：强制使用CPU运行（零修改，兼容性最好） #

在代码最开头添加配置，直接禁用TensorFlow对Metal GPU的调用，所有操作统一跑CPU，完全规避设备冲突：

import tensorflow as tf
tf.config.set_visible_devices([], 'GPU')

适用场景：对训练速度要求不高，不愿意修改现有代码逻辑

方案2：剥离训练相关节点，仅继承预训练权重后重新编译 #

你只需要预训练模型的推理权重做微调，不需要保留原模型的优化器状态，可以重建模型结构继承权重，用原生Metal支持的优化器重新编译：

from tensorflow.keras.models import load_model, Model

# 加载原模型，不读取编译配置
pretrained_model = load_model('pretrained.h5', compile=False)
# 基于原模型的输入输出重建纯推理结构，自动继承所有权重
finetune_model = Model(inputs=pretrained_model.input, outputs=pretrained_model.output)
# 用Metal原生支持的优化器、自定义损失指标重新编译
finetune_model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),
    loss="你的损失函数名",
    metrics=["你需要的评估指标"]
)
# 正常调用fit即可使用GPU训练
finetune_model.fit(train, steps_per_epoch=len(train), epochs=num_epochs, shuffle=True)

适用场景：希望保留GPU训练速度，可自行指定编译参数

方案3：开启软设备分配（保留GPU速度，无需修改模型） #

在代码开头添加配置，允许TensorFlow自动调度不支持GPU的操作到CPU运行，自动解决设备冲突：

import tensorflow as tf
tf.config.set_soft_device_placement(True)

适用场景：方案2无效时可以尝试，无需修改模型结构和编译逻辑

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