编写模型

如果您尝试做一些全新的事情,您可能希望从头开始实现一个模型。但是,在很多情况下,您可能对修改或扩展现有模型的某些组件感兴趣。因此,我们也提供了一些机制,让用户可以覆盖标准模型某些内部组件的行为。

注册新组件

对于用户经常想要自定义的常见概念,例如“骨干特征提取器”、“框头”,我们提供了一个注册机制,允许用户注入自定义实现,这些实现将立即可以在配置文件中使用。

例如,要添加一个新的骨干,在您的代码中导入以下代码

from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec

@BACKBONE_REGISTRY.register()
class ToyBackbone(Backbone):
  def __init__(self, cfg, input_shape):
    super().__init__()
    # create your own backbone
    self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=16, padding=3)

  def forward(self, image):
    return {"conv1": self.conv1(image)}

  def output_shape(self):
    return {"conv1": ShapeSpec(channels=64, stride=16)}

在这段代码中,我们实现了一个新的骨干,它遵循 Backbone 类接口,并将其注册到 BACKBONE_REGISTRY 中,该注册表需要 Backbone 的子类。导入这段代码后,detectron2 可以将类的名称链接到它的实现。因此,您可以编写以下代码

cfg = ...   # read a config
cfg.MODEL.BACKBONE.NAME = 'ToyBackbone'   # or set it in the config file
model = build_model(cfg)  # it will find `ToyBackbone` defined above

另一个例子,要为广义 R-CNN 元架构中的 ROI 头添加新的功能,您可以实现一个新的 ROIHeads 子类,并将其放在 ROI_HEADS_REGISTRY 中。 DensePoseMeshRCNN 是两个例子,它们实现了新的 ROIHeads 来执行新的任务。而且 projects/ 包含更多实现不同架构的例子。

可以在 API 文档 中找到完整的注册表列表。您可以在这些注册表中注册组件,以自定义模型的不同部分,或整个模型。

使用显式参数构建模型

注册表是连接配置文件中的名称与实际代码的桥梁。它们旨在涵盖用户经常需要替换的几个主要组件。但是,基于文本的配置文件的功能有时会受到限制,而一些更深入的自定义可能只能通过编写代码才能实现。

detectron2 中的大多数模型组件都有一个清晰的 __init__ 接口,该接口记录了它需要什么输入参数。使用自定义参数调用它们将为您提供模型的自定义变体。

例如,要在 Faster R-CNN 的框头上使用**自定义损失函数**,我们可以执行以下操作

  1. 损失目前是在 FastRCNNOutputLayers 中计算的。我们需要实现它的变体或子类,使用自定义损失函数,命名为 MyRCNNOutput

  2. 使用 box_predictor=MyRCNNOutput() 参数调用 StandardROIHeads,而不是内置的 FastRCNNOutputLayers。如果所有其他参数都应保持不变,这可以通过使用 可配置的 __init__ 机制轻松实现

    roi_heads = StandardROIHeads(
      cfg, backbone.output_shape(),
      box_predictor=MyRCNNOutput(...)
    )
    
  3. (可选) 如果我们想从配置文件中启用这个新模型,则需要注册

    @ROI_HEADS_REGISTRY.register()
    class MyStandardROIHeads(StandardROIHeads):
      def __init__(self, cfg, input_shape):
        super().__init__(cfg, input_shape,
                         box_predictor=MyRCNNOutput(...))