| 21 | |
| 22 | |
| 23 | class COAT(nn.Module): |
| 24 | def __init__(self, cfg): |
| 25 | super(COAT, self).__init__() |
| 26 | |
| 27 | backbone, _ = build_resnet(name="resnet50", pretrained=True) |
| 28 | anchor_generator = AnchorGenerator( |
| 29 | sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),) |
| 30 | ) |
| 31 | head = RPNHead( |
| 32 | in_channels=backbone.out_channels, |
| 33 | num_anchors=anchor_generator.num_anchors_per_location()[0], |
| 34 | ) |
| 35 | pre_nms_top_n = dict( |
| 36 | training=cfg.MODEL.RPN.PRE_NMS_TOPN_TRAIN, testing=cfg.MODEL.RPN.PRE_NMS_TOPN_TEST |
| 37 | ) |
| 38 | post_nms_top_n = dict( |
| 39 | training=cfg.MODEL.RPN.POST_NMS_TOPN_TRAIN, testing=cfg.MODEL.RPN.POST_NMS_TOPN_TEST |
| 40 | ) |
| 41 | rpn = RegionProposalNetwork( |
| 42 | anchor_generator=anchor_generator, |
| 43 | head=head, |
| 44 | fg_iou_thresh=cfg.MODEL.RPN.POS_THRESH_TRAIN, |
| 45 | bg_iou_thresh=cfg.MODEL.RPN.NEG_THRESH_TRAIN, |
| 46 | batch_size_per_image=cfg.MODEL.RPN.BATCH_SIZE_TRAIN, |
| 47 | positive_fraction=cfg.MODEL.RPN.POS_FRAC_TRAIN, |
| 48 | pre_nms_top_n=pre_nms_top_n, |
| 49 | post_nms_top_n=post_nms_top_n, |
| 50 | nms_thresh=cfg.MODEL.RPN.NMS_THRESH, |
| 51 | ) |
| 52 | |
| 53 | box_head = TransformerHead( |
| 54 | cfg=cfg, |
| 55 | trans_names=cfg.MODEL.TRANSFORMER.NAMES_1ST, |
| 56 | kernel_size=cfg.MODEL.TRANSFORMER.KERNEL_SIZE_1ST, |
| 57 | use_feature_mask=cfg.MODEL.TRANSFORMER.USE_MASK_1ST, |
| 58 | ) |
| 59 | box_head_2nd = TransformerHead( |
| 60 | cfg=cfg, |
| 61 | trans_names=cfg.MODEL.TRANSFORMER.NAMES_2ND, |
| 62 | kernel_size=cfg.MODEL.TRANSFORMER.KERNEL_SIZE_2ND, |
| 63 | use_feature_mask=cfg.MODEL.TRANSFORMER.USE_MASK_2ND, |
| 64 | ) |
| 65 | box_head_3rd = TransformerHead( |
| 66 | cfg=cfg, |
| 67 | trans_names=cfg.MODEL.TRANSFORMER.NAMES_3RD, |
| 68 | kernel_size=cfg.MODEL.TRANSFORMER.KERNEL_SIZE_3RD, |
| 69 | use_feature_mask=cfg.MODEL.TRANSFORMER.USE_MASK_3RD, |
| 70 | ) |
| 71 | |
| 72 | faster_rcnn_predictor = FastRCNNPredictor(2048, 2) |
| 73 | box_roi_pool = MultiScaleRoIAlign( |
| 74 | featmap_names=["feat_res4"], output_size=14, sampling_ratio=2 |
| 75 | ) |
| 76 | box_predictor = BBoxRegressor(2048, num_classes=2, bn_neck=cfg.MODEL.ROI_HEAD.BN_NECK) |
| 77 | roi_heads = CascadedROIHeads( |
| 78 | cfg=cfg, |
| 79 | # Cascade Transformer Head |
| 80 | faster_rcnn_predictor=faster_rcnn_predictor, |