如何在 Detectron2 中设置自定义类权重答案

【问题标题】：How set custom class weights in Detectron2如何在 Detectron2 中设置自定义类权重
【发布时间】：2020-10-20 17:41:34
【问题描述】：

我将https://detectron2.readthedocs.io/tutorials/install.html 用于包含其他类和对象的数据集。

而且我的数据集不平衡。我希望为每个班级设置不同的权重。我该怎么做？

【问题讨论】：

你可以检查文档中的扩充，detectron2.readthedocs.io/modules/…，或者如果这不起作用，在将它们输入数据加载器之前自己翻转那些代表性不足的类？

标签： python machine-learning computer-vision pytorch

【解决方案1】：

很遗憾，如果没有writing your own components，则无法配置此功能。

快速完成此操作的一种方法是编写一个新的头，该头继承自包含您的损失的头。然后，损失将被替换为使用您的损失权重初始化的新损失对象。

以DeepLabV3+ 为例，如下所示：

import torch
from torch import nn
from detectron2.modeling import SEM_SEG_HEADS_REGISTRY
from detectron2.projects.deeplab import DeepLabCE, DeepLabV3PlusHead

@SEM_SEG_HEADS_REGISTRY.register()
class MyNewHead(DeepLabV3PlusHead):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        weight = torch.Tensor([0.4, 0.6])  # Adapt to your case
        if self.loss_type == "cross_entropy":
            self.loss = nn.CrossEntropyLoss(
                reduction="mean", ignore_index=self.ignore_value, weight=weight
            )
        elif self.loss_type == "hard_pixel_mining":
            self.loss = DeepLabCE(
                ignore_label=self.ignore_value,
                top_k_percent_pixels=0.2,
                weight=weight,
            )
        else:
            raise ValueError("Unexpected loss type: %s" % self.loss_type)

然后，修改配置文件以选择新头：

MODEL:  
  SEM_SEG_HEAD:
    NAME: "MyNewHead"

【讨论】：