Detectron2在AzureML中训练与验证损失日志记录问题

我在AzureML中训练Detectron2模型，AzureML支持指标日志功能。Detectron2默认会记录总损失、分类器损失、边界框损失等，但我不清楚这些损失是训练损失还是验证损失，也不了解它的防过拟合机制（是否会保存验证损失最低时的权重）。

我尝试通过自定义Hook来记录训练和验证损失，但当前实现的代码得到的损失数值和Detectron2默认日志不符，比如第99次迭代我计算的总损失是1.936，默认日志显示为2.097。我想咨询以下问题：

1. 如何正确记录训练与验证损失指标？
2. Detectron2实际的损失计算方式是什么？
3. 它会保存验证损失最低时的权重，还是仅在迭代结束后保存？

我的当前实现代码如下：

# After setting up the cfg

from detectron2.engine import HookBase
from detectron2.data import build_detection_train_loader
import detectron2.utils.comm as comm

# Test/vali loss
from detectron2.utils.events import get_event_storage

class TrainingLoss(HookBase):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg.clone()
        self.cfg.DATASETS.TRAIN = self.cfg.DATASETS.TRAIN
        self._loader = iter(build_detection_train_loader(self.cfg))

    def after_step(self):
        data = next(self._loader)
        with torch.no_grad():
            loss_dict = self.trainer.model(data)

            losses = sum(loss_dict.values())
            assert torch.isfinite(losses).all(), loss_dict

            loss_dict_reduced = {"val_" + k: v.item() for k, v in
                                 comm.reduce_dict(loss_dict).items()}
            losses_reduced = sum(loss for loss in loss_dict_reduced.values())
            if comm.is_main_process():
                self.trainer.storage.put_scalars(total_val_loss=losses_reduced,
                                                 **loss_dict_reduced)

            print(f"Training Loss (Iteration {self.trainer.iter}): {losses_reduced}")

class ValidationLoss(HookBase):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg.clone()
        self.cfg.DATASETS.TRAIN = cfg.DATASETS.TEST
        self._loader = iter(build_detection_train_loader(self.cfg))

    def after_step(self):
        data = next(self._loader)
        with torch.no_grad():
            loss_dict = self.trainer.model(data)

            losses = sum(loss_dict.values())
            assert torch.isfinite(losses).all(), loss_dict

            loss_dict_reduced = {"val_" + k: v.item() for k, v in
                                 comm.reduce_dict(loss_dict).items()}
            losses_reduced = sum(loss for loss in loss_dict_reduced.values())
            if comm.is_main_process():
                self.trainer.storage.put_scalars(total_val_loss=losses_reduced,
                                                 **loss_dict_reduced)

            print(f"Vali Loss (Iteration {self.trainer.iter}): {losses_reduced}")


trainer = DefaultTrainer(cfg)
val_loss = ValidationLoss(cfg)
train_loss = TrainingLoss(cfg)
trainer.register_hooks([val_loss])
trainer.register_hooks([train_loss])
trainer.resume_or_load(resume=False)
trainer.train()

问题解答 #

1. 如何正确记录训练与验证损失指标 #

Detectron2默认日志里的损失都是训练损失，是当前迭代批次计算出的、经过分布式归约后的损失值。你的自定义Hook存在几个核心问题，导致数值不符：

• TrainingLoss类完全冗余：DefaultTrainer本身已经在训练步骤中自动记录训练损失，不需要额外加载训练集重复计算。
• ValidationLoss实现错误：用build_detection_train_loader加载验证集会应用训练数据增强（随机翻转、裁剪等），导致验证损失失真；且每次step只取单个batch计算，得到的不是验证集平均损失，和默认训练损失（单批次）不具备可比性。

正确的实现方式：

• 不要在每个step后计算验证损失，按固定间隔（比如每500次迭代或每个epoch）计算整个验证集的平均损失，避免拖慢训练。
• 使用build_detection_test_loader加载验证集，跳过训练数据增强。
• 遍历验证集所有batch，计算平均损失后再记录到AzureML。

修正后的验证损失Hook示例：

class ValidationLoss(HookBase):
    def __init__(self, cfg, eval_period=500):
        super().__init__()
        self.cfg = cfg.clone()
        self._loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0])
        self._eval_period = eval_period

    def after_step(self):
        # 仅在指定迭代间隔或最后一次迭代计算验证损失
        if self.trainer.iter % self._eval_period != 0 and self.trainer.iter != self.trainer.max_iter - 1:
            return
        
        total_loss = 0.0
        loss_sum = {}
        with torch.no_grad():
            for data in self._loader:
                loss_dict = self.trainer.model(data)
                # 分布式环境下归约损失
                loss_dict_reduced = comm.reduce_dict(loss_dict)
                total_loss += sum(loss_dict_reduced.values()).item()
                for k, v in loss_dict_reduced.items():
                    loss_sum[k] = loss_sum.get(k, 0) + v.item()
        
        # 计算验证集平均损失
        avg_total_loss = total_loss / len(self._loader)
        avg_loss_dict = {k: v / len(self._loader) for k, v in loss_sum.items()}
        
        if comm.is_main_process():
            self.trainer.storage.put_scalars(
                total_val_loss=avg_total_loss,
                **{f"val_{k}": v for k, v in avg_loss_dict.items()}
            )
            print(f"Validation Loss (Iteration {self.trainer.iter}): {avg_total_loss}")

训练损失无需额外Hook，DefaultTrainer会自动将total_loss、loss_cls、loss_box_reg等训练损失指标记录到日志，可直接在AzureML中查看。

2. Detectron2实际的损失计算方式 #

Detectron2的损失是多任务加权损失和，具体流程：

1. 模型前向传播时，针对分类、边界框回归、掩码预测等不同任务，分别计算对应损失（比如分类用交叉熵，框回归用Smooth L1）。
2. 每个损失项会乘以配置中对应的权重系数（比如MODEL.ROI_HEADS.LOSS_WEIGHTS里的loss_box_reg默认权重为5.0）。
3. 所有加权后的损失相加得到总损失，用于反向传播更新参数。
4. 分布式训练中，会对所有GPU的损失值进行归约（默认是平均），再记录到日志。

你自定义代码中直接求和未加权的损失，且没有对齐分布式归约逻辑，这是数值和默认日志不符的核心原因。

3. 权重保存机制 #

DefaultTrainer默认行为：

• 每SOLVER.CHECKPOINT_PERIOD次迭代保存一次权重（默认5000次迭代）。
• 训练结束时保存最终权重。
• 默认不会自动保存验证损失最低的权重，需自定义逻辑实现。

若要保存最优权重，可在ValidationLoss Hook中添加判断逻辑：

class ValidationLoss(HookBase):
    def __init__(self, cfg, eval_period=500):
        super().__init__()
        self.cfg = cfg.clone()
        self._loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0])
        self._eval_period = eval_period
        self.best_val_loss = float("inf")  # 初始化最优损失为无穷大

    def after_step(self):
        if self.trainer.iter % self._eval_period != 0 and self.trainer.iter != self.trainer.max_iter - 1:
            return
        
        # ... 计算平均验证损失的代码 ...
        
        if comm.is_main_process():
            self.trainer.storage.put_scalars(
                total_val_loss=avg_total_loss,
                **{f"val_{k}": v for k, v in avg_loss_dict.items()}
            )
            print(f"Validation Loss (Iteration {self.trainer.iter}): {avg_total_loss}")
            
            # 保存验证损失最低的权重
            if avg_total_loss < self.best_val_loss:
                self.best_val_loss = avg_total_loss
                self.trainer.save(f"model_best_{self.trainer.iter}")

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