| 27 | |
| 28 | |
| 29 | class DynamicHead(nn.Module): |
| 30 | |
| 31 | def __init__(self, cfg, roi_input_shape): |
| 32 | super().__init__() |
| 33 | |
| 34 | # Build RoI. |
| 35 | box_pooler = self._init_box_pooler(cfg, roi_input_shape) |
| 36 | self.box_pooler = box_pooler |
| 37 | |
| 38 | # Build heads. |
| 39 | num_classes = cfg.MODEL.SparseRCNN.NUM_CLASSES |
| 40 | d_model = cfg.MODEL.SparseRCNN.HIDDEN_DIM |
| 41 | dim_feedforward = cfg.MODEL.SparseRCNN.DIM_FEEDFORWARD |
| 42 | nhead = cfg.MODEL.SparseRCNN.NHEADS |
| 43 | dropout = cfg.MODEL.SparseRCNN.DROPOUT |
| 44 | activation = cfg.MODEL.SparseRCNN.ACTIVATION |
| 45 | num_heads = cfg.MODEL.SparseRCNN.NUM_HEADS |
| 46 | rcnn_head = RCNNHead(cfg, d_model, num_classes, dim_feedforward, nhead, dropout, activation) |
| 47 | self.head_series = _get_clones(rcnn_head, num_heads) |
| 48 | self.return_intermediate = cfg.MODEL.SparseRCNN.DEEP_SUPERVISION |
| 49 | |
| 50 | # Init parameters. |
| 51 | self.use_focal = cfg.MODEL.SparseRCNN.USE_FOCAL |
| 52 | self.num_classes = num_classes |
| 53 | if self.use_focal: |
| 54 | prior_prob = cfg.MODEL.SparseRCNN.PRIOR_PROB |
| 55 | self.bias_value = -math.log((1 - prior_prob) / prior_prob) |
| 56 | self._reset_parameters() |
| 57 | |
| 58 | def _reset_parameters(self): |
| 59 | # init all parameters. |
| 60 | for p in self.parameters(): |
| 61 | if p.dim() > 1: |
| 62 | nn.init.xavier_uniform_(p) |
| 63 | |
| 64 | # initialize the bias for focal loss. |
| 65 | if self.use_focal: |
| 66 | if p.shape[-1] == self.num_classes: |
| 67 | nn.init.constant_(p, self.bias_value) |
| 68 | |
| 69 | @staticmethod |
| 70 | def _init_box_pooler(cfg, input_shape): |
| 71 | |
| 72 | in_features = cfg.MODEL.ROI_HEADS.IN_FEATURES |
| 73 | pooler_resolution = cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION |
| 74 | pooler_scales = tuple(1.0 / input_shape[k].stride for k in in_features) |
| 75 | sampling_ratio = cfg.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO |
| 76 | pooler_type = cfg.MODEL.ROI_BOX_HEAD.POOLER_TYPE |
| 77 | |
| 78 | # If StandardROIHeads is applied on multiple feature maps (as in FPN), |
| 79 | # then we share the same predictors and therefore the channel counts must be the same |
| 80 | in_channels = [input_shape[f].channels for f in in_features] |
| 81 | # Check all channel counts are equal |
| 82 | assert len(set(in_channels)) == 1, in_channels |
| 83 | |
| 84 | box_pooler = ROIPooler( |
| 85 | output_size=pooler_resolution, |
| 86 | scales=pooler_scales, |