【OWOD论文】开放世界中OD代码_2

【OWOD论文】开放世界中OD代码_2_模型部分

news2024/11/15 11:49:08

简介

本文记录OWOD代码中的模型代码部分。数据部分可看我上一个博客【【OWOD论文】开放世界中OD代码_1_数据部分-CSDN博客】

模型代码

1 起步

在代码中找到 detectron2\engine\defaults.py DefaultTrainer类 __init__方法

根据上述 build_model 回溯到 detectron2\modeling\meta_arch\build.py 时，发现仅有下面的内容，明显 META_ARCH_REGISTRY 是关键：

根据代码，可发现 META_ARCH_REGISTRY 定义为

from detectron2.utils.registry import Registry

META_ARCH_REGISTRY = Registry("META_ARCH")  # noqa F401 isort:skip
META_ARCH_REGISTRY.__doc__ = """

由此找到 detectron2/utils/registry.py 中的 Registry 可发现下面的代码：

# Keep this module for backward compatibility.
from fvcore.common.registry import Registry  # noqa

__all__ = ["Registry"]

貌似还是没找到想要的，其实 detectron2 使用了 Registry 注解用于模型的注入。由此我们直接招到对应的模型部分，因为OWOD基于Faster RCNN，所以直接在 detectron2\modeling\meta_arch 中找到 rcnn.py

所以模型真正意义上在这里定义的，基于注解，将模型注入。

2 具体代码

GeneralizedRCNN的代码逻辑如下：
Generalized R-CNN. Any models that contains the following three components:
1. Per-image feature extraction (aka backbone)
2. Region proposal generation
3. Per-region feature extraction and prediction

当模型加载时，并不是直接执行 __init__ 方法，注意到 GeneralizadRCNN 中有个类方法，使用@classmethod 注解标明，方法初始化时，先加载这个静态方法进行了变量初始化。然后才是__init__ 方法。下面方法中的 cls 是个类占位符变量，这个可以看下 python 中的类注解方法定义。

@classmethod
def from_config(cls, cfg):
    backbone = build_backbone(cfg)
    return {
            "backbone": backbone,
            "proposal_generator": build_proposal_generator(cfg, backbone.output_shape()),
            "roi_heads": build_roi_heads(cfg, backbone.output_shape()),
            "input_format": cfg.INPUT.FORMAT,
            "vis_period": cfg.VIS_PERIOD,
            "pixel_mean": cfg.MODEL.PIXEL_MEAN,
            "pixel_std": cfg.MODEL.PIXEL_STD,
        }

这里面有这两个方法，一个是生成 Proposal boxes，一个为 roi_heads

"proposal_generator": build_proposal_generator(cfg, backbone.output_shape()),
"roi_heads": build_roi_heads(cfg, backbone.output_shape()),

2.1 build_proposal_generator

先来看 build_proposal_generator 部分，本质上为Faster RCNN的建议框生成过程，其初始化的代码部分为

#cfg 一些配置文件
    #input_shape:{"res4":[1024,16]}
    #即 input_shape 为RPN.IN_FEATURES 
    #定义的Resnet输出的feature
    #该代码中使用的是仅输出 ['res4']    
    @classmethod
    def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
        #res4
        in_features = cfg.MODEL.RPN.IN_FEATURES
        ret = {
            "in_features": in_features,
            #0
            "min_box_size": cfg.MODEL.PROPOSAL_GENERATOR.MIN_SIZE,
            #0.7
            "nms_thresh": cfg.MODEL.RPN.NMS_THRESH,
            #256
            "batch_size_per_image": cfg.MODEL.RPN.BATCH_SIZE_PER_IMAGE,
            #0.5
            "positive_fraction": cfg.MODEL.RPN.POSITIVE_FRACTION,
            "loss_weight": {
                #loss_rpn_cls 1.0 loss_rpn_loc:1.0*1.0
                "loss_rpn_cls": cfg.MODEL.RPN.LOSS_WEIGHT,
                "loss_rpn_loc": cfg.MODEL.RPN.BBOX_REG_LOSS_WEIGHT * cfg.MODEL.RPN.LOSS_WEIGHT,
            },
            #-1 
            #Remove RPN anchors that 
            #go outside the image by BOUNDARY_THRESH pixels
            "anchor_boundary_thresh": cfg.MODEL.RPN.BOUNDARY_THRESH,
            "box2box_transform": Box2BoxTransform(weights=cfg.MODEL.RPN.BBOX_REG_WEIGHTS),
            #smooth_l1
            "box_reg_loss_type": cfg.MODEL.RPN.BBOX_REG_LOSS_TYPE,
            #0.0
            "smooth_l1_beta": cfg.MODEL.RPN.SMOOTH_L1_BETA,
        }
        #12000 / 6000
        ret["pre_nms_topk"] = (cfg.MODEL.RPN.PRE_NMS_TOPK_TRAIN, cfg.MODEL.RPN.PRE_NMS_TOPK_TEST)
        #2000 / 1000
        ret["post_nms_topk"] = (cfg.MODEL.RPN.POST_NMS_TOPK_TRAIN, cfg.MODEL.RPN.POST_NMS_TOPK_TEST)

        ret["anchor_generator"] = build_anchor_generator(cfg, [input_shape[f] for f in in_features])
        ret["anchor_matcher"] = Matcher(
            cfg.MODEL.RPN.IOU_THRESHOLDS, cfg.MODEL.RPN.IOU_LABELS, allow_low_quality_matches=True
        )
        ret["head"] = build_rpn_head(cfg, [input_shape[f] for f in in_features])
        return ret

其中需要关注的是下面的部分

ret["anchor_generator"] = build_anchor_generator(cfg, [input_shape[f] for f in in_features])
ret["anchor_matcher"] = Matcher(cfg.MODEL.RPN.IOU_THRESHOLDS, cfg.MODEL.RPN.IOU_LABELS, allow_low_quality_matches=True)
ret["head"] = build_rpn_head(cfg, [input_shape[f] for f in in_features])

build_anchor_generator 用于生成锚框，主要看 detectron2\modeling\anchor_generator.py 中的 DefaultAnchorGenerator 其 from_config 内容为

@classmethod
def from_config(cls, cfg, input_shape: List[ShapeSpec]):
     return {
          #[[32, 64, 128, 256, 512]]
          "sizes": cfg.MODEL.ANCHOR_GENERATOR.SIZES,
          #[[0.5, 1.0, 2.0]]
          "aspect_ratios": cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS,
          #[16]
          "strides": [x.stride for x in input_shape],
          #0.0
          "offset": cfg.MODEL.ANCHOR_GENERATOR.OFFSET,
        }

Matcher 用于匹配正负和忽略样本，具体代码参看 detectron2\modeling\Matcher.py