YOLOv3 Keras模型加载报错：yolo_head未定义问题求助

核心原因 #

使用save_weights_only=False保存完整模型时，Keras会将模型结构（包括自定义层/函数的调用关系）一并保存。加载时不仅需要指定损失函数yolo_loss，还需要包含模型结构中用到的所有自定义函数——yolo_head是yolo_body构建模型时直接调用的函数，必须加入custom_objects字典。

解决方案 #

方案1：显式指定所有用到的自定义函数 #

从model.py中导入所有模型结构依赖的自定义函数，加载时全部传入custom_objects：

from keras.models import load_model
from model import yolo_loss, yolo_head, yolo_boxes_and_scores, yolo_body

# 加载剪枝后的模型
pruned_model = load_model(
    'path/to/your/pruned_model.h5',
    custom_objects={
        'yolo_loss': yolo_loss,
        'yolo_head': yolo_head,
        'yolo_boxes_and_scores': yolo_boxes_and_scores
    }
)

方案2：直接导入整个model模块的字典（更简便） #

如果不确定哪些函数需要导入，可以直接将model模块的全局字典传入custom_objects，自动覆盖所有模块内的自定义函数：

from keras.models import load_model
import model

pruned_model = load_model(
    'path/to/your/pruned_model.h5',
    custom_objects=model.__dict__
)

额外检查项 #

• 确保加载模型的环境与保存时一致：Python3.6.5、Tensorflow1.13、对应版本的Keras（建议使用Keras 2.2.4，适配TF1.13）。
• 确认剪枝过程未破坏模型结构中对yolo_head等函数的调用逻辑，避免剪枝工具误删相关节点引用。

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