如何保存Mask RCNN训练生成的模型权重以用于后续训练使用？

场景1：同任务中断续训（数据集、类别数和之前训练完全一致） #

你遇到的问题核心是：加载公开预训练权重时加exclude参数是因为公开预训练的输出层维度适配COCO数据集，和你的自定义任务维度不匹配，需要排除后随机初始化输出层。但你自己训练的权重输出层维度和当前任务完全匹配，加exclude会把已训好的输出层权重丢弃，用随机值替代，大概率会出现训练不动、loss异常的问题。
操作步骤：

• 加载权重时去掉exclude参数，代码调整为：

model.load_weights('mask_rcnn_building_cfg_0001.h5', by_name=True)

• 训练启动时将起始epoch设置为权重对应的编号，比如加载的是第1轮的权重，起始epoch设为2，保证学习率调度器、日志记录等逻辑正常运行。

场景2：将训练好的权重作为通用预训练权重，用于其他不同类别数的任务 #

如果需要输出的权重可以像公开预训练权重一样适配不同任务，有两种实现方式：

方式1：保存权重时直接排除输出层 #

自定义回调函数，保存时过滤掉头部输出层的权重，参考代码：

class CustomCheckpoint(keras.callbacks.Callback):
    def __init__(self, save_path, exclude_layers):
        super().__init__()
        self.save_path = save_path
        self.exclude_layers = exclude_layers
    
    def on_epoch_end(self, epoch, logs=None):
        file_path = self.save_path.format(epoch=epoch+1)
        save_weights = []
        for layer in self.model.layers:
            if layer.name not in self.exclude_layers:
                save_weights.extend(layer.weights)
        self.model.save_weights(file_path, save_weights)

替换原回调配置：

# 要排除的输出层列表
exclude_head_layers = ["mrcnn_class_logits", "mrcnn_bbox_fc",  "mrcnn_bbox", "mrcnn_mask"]
callbacks = [
    keras.callbacks.TensorBoard(log_dir=self.log_dir,
                                histogram_freq=0, write_graph=True, write_images=False),
    CustomCheckpoint(self.checkpoint_path, exclude_head_layers),
]

注意调整self.checkpoint_path格式为带epoch占位符的形式，比如mask_rcnn_building_cfg_{epoch:04d}.h5，保证每轮生成对应编号的权重文件。

方式2：保留全量权重，加载时增加兼容参数 #

如果不需要改保存逻辑，给其他任务加载时增加skip_mismatch=True参数即可适配维度不匹配的问题，代码调整为：

model.load_weights('mask_rcnn_building_cfg_0001.h5', by_name=True, skip_mismatch=True, exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",  "mrcnn_bbox", "mrcnn_mask"])

常见问题排查 #

• 确认保存的权重文件没有损坏，可通过h5py工具打开验证权重参数是否正常写入
• 续训前检查模型各层的trainable属性设置是否符合预期，避免骨干层被冻结导致权重无更新
• 调整适配续训的学习率，不要使用从头训练的大学习率，避免权重被冲毁，也不要设置过小导致更新幅度过低

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