如何修改Detectron2中DefaultTrainer训练时的模型保存文件名

方法一：继承DefaultTrainer重写Checkpointer配置 #

DefaultTrainer的模型保存逻辑由DetectionCheckpointer处理，默认最终模型文件名为model_final.pth。你可以通过继承DefaultTrainer并重写build_checkpointer方法，指定自定义文件名：

from detectron2.engine import DefaultTrainer
from detectron2.checkpoint import DetectionCheckpointer

class CustomTrainer(DefaultTrainer):
    def build_checkpointer(self, cfg, model, optimizer, scheduler):
        return DetectionCheckpointer(
            model,
            cfg.OUTPUT_DIR,
            optimizer=optimizer,
            scheduler=scheduler,
            save_dir=cfg.OUTPUT_DIR,
            final_model_name="my_custom_model.pth"  # 此处替换为你的目标文件名
        )

# 使用自定义Trainer执行训练
cfg = get_cfg()
# ... 你的配置初始化代码 ...
trainer = CustomTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

方法二：训练完成后手动重命名文件 #

如果不想修改Trainer类结构，可以在训练结束后直接对默认生成的文件进行重命名，操作简单直接：

import os

cfg = get_cfg()
# ... 你的配置初始化代码 ...
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

# 重命名模型文件
default_model_path = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
target_model_path = os.path.join(cfg.OUTPUT_DIR, "my_custom_model.pth")

if os.path.exists(default_model_path):
    os.rename(default_model_path, target_model_path)
    print(f"模型已重命名为: {target_model_path}")

两种方法各有适用场景：方法一适合需要嵌入自动化训练流程的场景，从根源修改保存逻辑；方法二更适合快速调整、无需重构代码的临时需求。

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