FasterRCNN pytorch官方源码学习

基类 GeneralizedRCNN

class GeneralizedRCNN(nn.Module):
    def __init__(self, backbone, rpn, roi_heads, transform):
        super(GeneralizedRCNN, self).__init__()
        self.transform = transform
        self.backbone = backbone
        self.rpn = rpn
        self.roi_heads = roi_heads
        # used only on torchscript mode
        self._has_warned = False

    @torch.jit.unused
    def eager_outputs(self, losses, detections):
        # type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]
        if self.training:
            return losses

        return detections

    def forward(self, images, targets=None):
        if self.training and targets is None:
            raise ValueError("In training mode, targets should be passed")
        #将images的原始尺寸保存到original_image_sizes中
        original_image_sizes = torch.jit.annotate(List[Tuple[int, int]], [])
        for img in images:
            #这个地方 输入的images已经转为了tensor格式 shape是 [c, h, w]
            val = img.shape[-2:]
            assert len(val) == 2
            original_image_sizes.append((val[0], val[1]))
        #tranform:将图像标准化，并缩放到固定大小
        images, targets = self.transform(images, targets)
        #进入主干网络，提取图像特征
        features = self.backbone(images.tensors)
        #这一步应该是修正一下，统一一下格式
        #因为当backbone使用了fpn的时候，会有多个feature_map，被存放到了OrderedDict中
        #而没有使用的时候，只会返回一个feature_map
        #因此当只有1个的时候，也将其存在OrderedDict中，统一格式
        if isinstance(features, torch.Tensor):
            features = OrderedDict([('0', features)])
        #rpn
        proposals, proposal_losses = self.rpn(images, features, targets)
        #roi_heads:roi_pooling + 分类
        detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
        #后处理，进行NMS 同时将 box 通过 original_images_size映射回原图
        detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)

        losses = {
   }
        losses.update(detector_losses)
        losses.update(proposal_losses)

        if torch.jit.is_scripting():
            if not self._has_warned:
                warnings.warn("RCNN always returns a (Losses, Detections) tuple in scripting")
                self._has_warned = True
            return (losses, detections)
        else:
            #训练模式下，只返回losses；推断模式下 只返回detections
            return self.eager_outputs(losses, detections)

FasterRCNN

class FasterRCNN(GeneralizedRCNN):

    def __init__(self, backbone, num_classes=None,
                 # transform parameters
                 min_size=800, max_size=1333,
                 image_mean=None, image_std=None,
                 # RPN parameters
                 rpn_anchor_generator=None, rpn_head=None,
                 rpn_pre_nms_top_n_train=2000, rpn_pre_nms_top_n_test=1000,
                 rpn_post_nms_top_n_train=2000, rpn_post_nms_top_n_test=1000,
                 rpn_nms_thresh=0.7,
                 rpn_fg_iou_thresh=0.7, rpn_bg_iou_thresh=0.3,
                 rpn_batch_size_per_image=256, rpn_positive_fraction=0.5,
                 # Box parameters
                 box_roi_pool=None, box_head=None, box_predictor=None,
                 box_score_thresh=0.05, box_nms_thresh=0.5, box_detections_per_img=100,
                 box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5,
                 box_batch_size_per_image=512, box_positive_fraction=0.25,
                 bbox_reg_weights=None):

        out_channels = backbone.out_channels
        
        #创建默认的anchor_generator
        if rpn_anchor_generator is None:
            anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
            #这个地方为啥要 * len(anchor_sizes)？？？
            aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
            rpn_anchor_generator = AnchorGenerator(
                anchor_sizes, aspect_ratios
            )
        #
        if rpn_head is None:
            rpn_head = RPNHead(
                out_channels, rpn_anchor_generator.num_anchors_