open more parameters for maskrcnn like fasterrcnn

docs/apis/models/detection.md

@@ -250,7 +250,7 @@ paddlex.det.FasterRCNN(num_classes=81, backbone='ResNet50', with_fpn=True, aspec
 > **参数**
 
 > > - **num_classes** (int): 包含了背景类的类别数。默认为81。
-> > - **backbone** (str): FasterRCNN的backbone网络，取值范围为['ResNet18', 'ResNet50', 'ResNet50_vd', 'ResNet101', 'ResNet101_vd', 'HRNet_W18']。默认为'ResNet50'。
+> > - **backbone** (str): FasterRCNN的backbone网络，取值范围为['ResNet18', 'ResNet50', 'ResNet50_vd', 'ResNet101', 'ResNet101_vd', 'HRNet_W18', 'ResNet50_vd_ssld']。默认为'ResNet50'。
 > > - **with_fpn** (bool): 是否使用FPN结构。默认为True。
 > > - **aspect_ratios** (list): 生成anchor高宽比的可选值。默认为[0.5, 1.0, 2.0]。
 > > - **anchor_sizes** (list): 生成anchor大小的可选值。默认为[32, 64, 128, 256, 512]。

docs/apis/models/instance_segmentation.md

@@ -3,7 +3,7 @@
 ## MaskRCNN
 
 ```python
-paddlex.det.MaskRCNN(num_classes=81, backbone='ResNet50', with_fpn=True, aspect_ratios=[0.5, 1.0, 2.0], anchor_sizes=[32, 64, 128, 256, 512], input_channel=3)
+paddlex.det.MaskRCNN(num_classes=81, backbone='ResNet50', with_fpn=True, aspect_ratios=[0.5, 1.0, 2.0], anchor_sizes=[32, 64, 128, 256, 512], with_dcn=False, rpn_cls_loss='SigmoidCrossEntropy', rpn_focal_loss_alpha=0.25, rpn_focal_loss_gamma=2, rcnn_bbox_loss='SmoothL1Loss', rcnn_nms='MultiClassNMS', keep_top_k=100, nms_threshold=0.5, score_threshold=0.05, softnms_sigma=0.5, bbox_assigner='BBoxAssigner', fpn_num_channels=256, input_channel=3, rpn_batch_size_per_im=256, rpn_fg_fraction=0.5, test_pre_nms_top_n=None, test_post_nms_top_n=1000)
 
 ```
 
@@ -12,11 +12,29 @@ paddlex.det.MaskRCNN(num_classes=81, backbone='ResNet50', with_fpn=True, aspect_
 > **参数**
 
 > > - **num_classes** (int): 包含了背景类的类别数。默认为81。
-> > - **backbone** (str): MaskRCNN的backbone网络，取值范围为['ResNet18', 'ResNet50', 'ResNet50_vd', 'ResNet101', 'ResNet101_vd', 'HRNet_W18']。默认为'ResNet50'。
+> > - **backbone** (str): MaskRCNN的backbone网络，取值范围为['ResNet18', 'ResNet50', 'ResNet50_vd', 'ResNet101', 'ResNet101_vd', 'HRNet_W18', 'ResNet50_vd_ssld']。默认为'ResNet50'。
 > > - **with_fpn** (bool): 是否使用FPN结构。默认为True。
 > > - **aspect_ratios** (list): 生成anchor高宽比的可选值。默认为[0.5, 1.0, 2.0]。
 > > - **anchor_sizes** (list): 生成anchor大小的可选值。默认为[32, 64, 128, 256, 512]。
+> > - **with_dcn** (bool): backbone网络中是否使用deformable convolution network v2。默认为False。
+> > - **rpn_cls_loss** (str): RPN部分的分类损失函数，取值范围为['SigmoidCrossEntropy', 'SigmoidFocalLoss']。当遇到模型误检了很多背景区域时，可以考虑使用'SigmoidFocalLoss'，并调整适合的`rpn_focal_loss_alpha`和`rpn_focal_loss_gamma`。默认为'SigmoidCrossEntropy'。
+> > - **rpn_focal_loss_alpha** (float)：当RPN的分类损失函数设置为'SigmoidFocalLoss'时，用于调整正样本和负样本的比例因子，默认为0.25。当PN的分类损失函数设置为'SigmoidCrossEntropy'时，`rpn_focal_loss_alpha`的设置不生效。
+> > - **rpn_focal_loss_gamma** (float): 当RPN的分类损失函数设置为'SigmoidFocalLoss'时，用于调整易分样本和难分样本的比例因子，默认为2。当RPN的分类损失函数设置为'SigmoidCrossEntropy'时，`rpn_focal_loss_gamma`的设置不生效。
+> > - **rcnn_bbox_loss** (str): RCNN部分的位置回归损失函数，取值范围为['SmoothL1Loss', 'CIoULoss']。默认为'SmoothL1Loss'。
+> > - **rcnn_nms** (str): RCNN部分的非极大值抑制的计算方法，取值范围为['MultiClassNMS', 'MultiClassSoftNMS','MultiClassCiouNMS']。默认为'MultiClassNMS'。当选择'MultiClassNMS'时，可以将`keep_top_k`设置成100、`nms_threshold`设置成0.5、`score_threshold`设置成0.05。当选择'MultiClassSoftNMS'时，可以将`keep_top_k`设置为300、`score_threshold`设置为0.01、`softnms_sigma`设置为0.5。当选择'MultiClassCiouNMS'时，可以将`keep_top_k`设置为100、`score_threshold`设置成0.05、`nms_threshold`设置成0.5。
+> > - **keep_top_k** (int): RCNN部分在进行非极大值抑制计算后，每张图像保留最多保存`keep_top_k`个检测框。默认为100。
+> > - **nms_threshold** (float): RCNN部分在进行非极大值抑制时，用于剔除检测框所需的IoU阈值。当`rcnn_nms`设置为`MultiClassSoftNMS`时，`nms_threshold`的设置不生效。默认为0.5。
+> > - **score_threshold** (float): RCNN部分在进行非极大值抑制前，用于过滤掉低置信度边界框所需的置信度阈值。默认为0.05。
+> > - **softnms_sigma** (float): 当`rcnn_nms`设置为`MultiClassSoftNMS`时，用于调整被抑制的检测框的置信度，调整公式为`score = score * weights, weights = exp(-(iou * iou) / softnms_sigma)`。默认设为0.5。
+> > - **bbox_assigner** (str): 训练阶段，RCNN部分生成正负样本的采样方式。可选范围为['BBoxAssigner', 'LibraBBoxAssigner']。当目标物体的区域只占原始图像的一小部分时，可以考虑采用[LibraRCNN](https://arxiv.org/abs/1904.02701)中提出的IoU-balanced Sampling采样方式来获取更多的难分负样本，设置为'LibraBBoxAssigner'即可。默认为'BBoxAssigner'。
+> > - **fpn_num_channels** (int): FPN部分特征层的通道数量。默认为256。
 > > - **input_channel** (int): 输入图像的通道数量。默认为3。
+> > - **rpn_batch_size_per_im** (int): 训练阶段，RPN部分每张图片的正负样本的数量总和。默认为256。
+> > - **rpn_fg_fraction** (float): 训练阶段，RPN部分每张图片的正负样本数量总和中正样本的占比。默认为0.5。
+> > - **test_pre_nms_top_n** (int)：预测阶段，RPN部分做非极大值抑制计算的候选框的数量。若设置为None, 有FPN结构的话，`test_pre_nms_top_n`会被设置成6000, 无FPN结构的话，`test_pre_nms_top_n`会被设置成1000。默认为None。
+> > - **test_post_nms_top_n** (int): 预测阶段，RPN部分做完非极大值抑制后保留的候选框的数量。默认为1000。
+
+
 
 #### train
 
@@ -74,7 +92,7 @@ evaluate(self, eval_dataset, batch_size=1, epoch_id=None, metric=None, return_de
 predict(self, img_file, transforms=None)
 ```
 
-> MaskRCNN模型预测接口。需要注意的是，只有在训练过程中定义了eval_dataset，模型在保存时才会将预测时的图像处理流程保存在`FasterRCNN.test_transforms`和`FasterRCNN.eval_transforms`中。如未在训练时定义eval_dataset，那在调用预测`predict`接口时，用户需要再重新定义test_transforms传入给`predict`接口。
+> MaskRCNN模型预测接口。需要注意的是，只有在训练过程中定义了eval_dataset，模型在保存时才会将预测时的图像处理流程保存在`MaskRCNN.test_transforms`和`MaskRCNN.eval_transforms`中。如未在训练时定义eval_dataset，那在调用预测`predict`接口时，用户需要再重新定义test_transforms传入给`predict`接口。
 
 > **参数**
 >
@@ -92,7 +110,7 @@ predict(self, img_file, transforms=None)
 batch_predict(self, img_file_list, transforms=None)
 ```
 
-> MaskRCNN模型批量预测接口。需要注意的是，只有在训练过程中定义了eval_dataset，模型在保存时才会将预测时的图像处理流程保存在`FasterRCNN.test_transforms`和`FasterRCNN.eval_transforms`中。如未在训练时定义eval_dataset，那在调用预测`batch_predict`接口时，用户需要再重新定义test_transforms传入给`batch_predict`接口。
+> MaskRCNN模型批量预测接口。需要注意的是，只有在训练过程中定义了eval_dataset，模型在保存时才会将预测时的图像处理流程保存在`MaskRCNN.test_transforms`和`MaskRCNN.eval_transforms`中。如未在训练时定义eval_dataset，那在调用预测`batch_predict`接口时，用户需要再重新定义test_transforms传入给`batch_predict`接口。
 
 > **参数**
 >

paddlex/cv/models/mask_rcnn.py

@@ -34,11 +34,43 @@ class MaskRCNN(FasterRCNN):
     Args:
         num_classes (int): 包含了背景类的类别数。默认为81。
         backbone (str): MaskRCNN的backbone网络，取值范围为['ResNet18', 'ResNet50',
-            'ResNet50_vd', 'ResNet101', 'ResNet101_vd', 'HRNet_W18']。默认为'ResNet50'。
+            'ResNet50_vd', 'ResNet101', 'ResNet101_vd', 'HRNet_W18', 'ResNet50_vd_ssld']。默认为'ResNet50'。
         with_fpn (bool): 是否使用FPN结构。默认为True。
         aspect_ratios (list): 生成anchor高宽比的可选值。默认为[0.5, 1.0, 2.0]。
         anchor_sizes (list): 生成anchor大小的可选值。默认为[32, 64, 128, 256, 512]。
+        with_dcn (bool): backbone网络中是否使用deformable convolution network v2。默认为False。
+        rpn_cls_loss (str): RPN部分的分类损失函数，取值范围为['SigmoidCrossEntropy', 'SigmoidFocalLoss']。
+            当遇到模型误检了很多背景区域时，可以考虑使用'SigmoidFocalLoss'，并调整适合的`rpn_focal_loss_alpha`
+            和`rpn_focal_loss_gamma`。默认为'SigmoidCrossEntropy'。
+        rpn_focal_loss_alpha (float)：当RPN的分类损失函数设置为'SigmoidFocalLoss'时，用于调整
+            正样本和负样本的比例因子，默认为0.25。当PN的分类损失函数设置为'SigmoidCrossEntropy'时，
+            `rpn_focal_loss_alpha`的设置不生效。
+        rpn_focal_loss_gamma (float): 当RPN的分类损失函数设置为'SigmoidFocalLoss'时，用于调整
+            易分样本和难分样本的比例因子，默认为2。当RPN的分类损失函数设置为'SigmoidCrossEntropy'时，
+            `rpn_focal_loss_gamma`的设置不生效。
+        rcnn_bbox_loss (str): RCNN部分的位置回归损失函数，取值范围为['SmoothL1Loss', 'CIoULoss']。
+            默认为'SmoothL1Loss'。
+        rcnn_nms (str): RCNN部分的非极大值抑制的计算方法，取值范围为['MultiClassNMS', 'MultiClassSoftNMS',
+            'MultiClassCiouNMS']。默认为'MultiClassNMS'。当选择'MultiClassNMS'时，可以将`keep_top_k`设置成100、
+            `nms_threshold`设置成0.5、`score_threshold`设置成0.05。当选择'MultiClassSoftNMS'时，可以将`keep_top_k`设置为300、
+            `score_threshold`设置为0.01、`softnms_sigma`设置为0.5。当选择'MultiClassCiouNMS'时，可以将`keep_top_k`设置为100、
+            `score_threshold`设置成0.05、`nms_threshold`设置成0.5。
+        keep_top_k (int): RCNN部分在进行非极大值抑制计算后，每张图像保留最多保存`keep_top_k`个检测框。默认为100。
+        nms_threshold (float): RCNN部分在进行非极大值抑制时，用于剔除检测框所需的IoU阈值。
+            当`rcnn_nms`设置为`MultiClassSoftNMS`时，`nms_threshold`的设置不生效。默认为0.5。
+        score_threshold (float): RCNN部分在进行非极大值抑制前，用于过滤掉低置信度边界框所需的置信度阈值。默认为0.05。
+        softnms_sigma (float): 当`rcnn_nms`设置为`MultiClassSoftNMS`时，用于调整被抑制的检测框的置信度，
+            调整公式为`score = score * weights, weights = exp(-(iou * iou) / softnms_sigma)`。默认设为0.5。
+        bbox_assigner (str): 训练阶段，RCNN部分生成正负样本的采样方式。可选范围为['BBoxAssigner', 'LibraBBoxAssigner']。
+            当目标物体的区域只占原始图像的一小部分时，使用`LibraBBoxAssigner`采样方式模型效果更佳。默认为'BBoxAssigner'。
+        fpn_num_channels (int): FPN部分特征层的通道数量。默认为256。
         input_channel (int): 输入图像的通道数量。默认为3。
+        rpn_batch_size_per_im (int): 训练阶段，RPN部分每张图片的正负样本的数量总和。默认为256。
+        rpn_fg_fraction (float): 训练阶段，RPN部分每张图片的正负样本数量总和中正样本的占比。默认为0.5。
+        test_pre_nms_top_n (int)：预测阶段，RPN部分做非极大值抑制计算的候选框的数量。若设置为None,
+            有FPN结构的话，`test_pre_nms_top_n`会被设置成6000, 无FPN结构的话，`test_pre_nms_top_n`会被设置成
+            1000。默认为None。
+        test_post_nms_top_n (int): 预测阶段，RPN部分做完非极大值抑制后保留的候选框的数量。默认为1000。
     """
 
     def __init__(self,
@@ -47,11 +79,27 @@ class MaskRCNN(FasterRCNN):
                  with_fpn=True,
                  aspect_ratios=[0.5, 1.0, 2.0],
                  anchor_sizes=[32, 64, 128, 256, 512],
-                 input_channel=3):
+                 with_dcn=False,
+                 rpn_cls_loss='SigmoidCrossEntropy',
+                 rpn_focal_loss_alpha=0.25,
+                 rpn_focal_loss_gamma=2,
+                 rcnn_bbox_loss='SmoothL1Loss',
+                 rcnn_nms='MultiClassNMS',
+                 keep_top_k=100,
+                 nms_threshold=0.5,
+                 score_threshold=0.05,
+                 softnms_sigma=0.5,
+                 bbox_assigner='BBoxAssigner',
+                 fpn_num_channels=256,
+                 input_channel=3,
+                 rpn_batch_size_per_im=256,
+                 rpn_fg_fraction=0.5,
+                 test_pre_nms_top_n=None,
+                 test_post_nms_top_n=1000):
         self.init_params = locals()
         backbones = [
             'ResNet18', 'ResNet50', 'ResNet50_vd', 'ResNet101', 'ResNet101_vd',
-            'HRNet_W18'
+            'HRNet_W18', 'ResNet50_vd_ssld'
         ]
         assert backbone in backbones, "backbone should be one of {}".format(
             backbones)
@@ -59,6 +107,7 @@ class MaskRCNN(FasterRCNN):
         self.backbone = backbone
         self.num_classes = num_classes
         self.with_fpn = with_fpn
+        self.aspect_ratios = aspect_ratios
         self.anchor_sizes = anchor_sizes
         self.labels = None
         if with_fpn:
@@ -66,24 +115,60 @@ class MaskRCNN(FasterRCNN):
         else:
             self.mask_head_resolution = 14
         self.fixed_input_shape = None
+        self.with_dcn = with_dcn
+        rpn_cls_losses = ['SigmoidFocalLoss', 'SigmoidCrossEntropy']
+        assert rpn_cls_loss in rpn_cls_losses, "rpn_cls_loss should be one of {}".format(
+            rpn_cls_losses)
+        self.rpn_cls_loss = rpn_cls_loss
+        self.rpn_focal_loss_alpha = rpn_focal_loss_alpha
+        self.rpn_focal_loss_gamma = rpn_focal_loss_gamma
+        self.rcnn_bbox_loss = rcnn_bbox_loss
+        self.rcnn_nms = rcnn_nms
+        self.keep_top_k = keep_top_k
+        self.nms_threshold = nms_threshold
+        self.score_threshold = score_threshold
+        self.softnms_sigma = softnms_sigma
+        self.bbox_assigner = bbox_assigner
+        self.fpn_num_channels = fpn_num_channels
         self.input_channel = input_channel
-        self.with_dcn = False
+        self.rpn_batch_size_per_im = rpn_batch_size_per_im
+        self.rpn_fg_fraction = rpn_fg_fraction
+        self.test_pre_nms_top_n = test_pre_nms_top_n
+        self.test_post_nms_top_n = test_post_nms_top_n
 
     def build_net(self, mode='train'):
         train_pre_nms_top_n = 2000 if self.with_fpn else 12000
         test_pre_nms_top_n = 1000 if self.with_fpn else 6000
+        if self.test_pre_nms_top_n is not None:
+            test_pre_nms_top_n = self.test_pre_nms_top_n
         num_convs = 4 if self.with_fpn else 0
         model = paddlex.cv.nets.detection.MaskRCNN(
             backbone=self._get_backbone(self.backbone),
             num_classes=self.num_classes,
             mode=mode,
             with_fpn=self.with_fpn,
+            aspect_ratios=self.aspect_ratios,
+            anchor_sizes=self.anchor_sizes,
             train_pre_nms_top_n=train_pre_nms_top_n,
             test_pre_nms_top_n=test_pre_nms_top_n,
             num_convs=num_convs,
             mask_head_resolution=self.mask_head_resolution,
             fixed_input_shape=self.fixed_input_shape,
-            input_channel=self.input_channel)
+            rpn_cls_loss=self.rpn_cls_loss,
+            rpn_focal_loss_alpha=self.rpn_focal_loss_alpha,
+            rpn_focal_loss_gamma=self.rpn_focal_loss_gamma,
+            rcnn_bbox_loss=self.rcnn_bbox_loss,
+            rcnn_nms=self.rcnn_nms,
+            keep_top_k=self.keep_top_k,
+            nms_threshold=self.nms_threshold,
+            score_threshold=self.score_threshold,
+            softnms_sigma=self.softnms_sigma,
+            bbox_assigner=self.bbox_assigner,
+            fpn_num_channels=self.fpn_num_channels,
+            input_channel=self.input_channel,
+            rpn_batch_size_per_im=self.rpn_batch_size_per_im,
+            rpn_fg_fraction=self.rpn_fg_fraction,
+            test_post_nms_top_n=self.test_post_nms_top_n)
         inputs = model.generate_inputs()
         if mode == 'train':
             model_out = model.build_net(inputs)

paddlex/cv/nets/detection/mask_rcnn.py

@@ -27,6 +27,8 @@ from .roi_extractor import (RoIAlign, FPNRoIAlign)
 from .bbox_head import (BBoxHead, TwoFCHead)
 from .mask_head import MaskHead
 from ..resnet import ResNetC5
+from .loss.diou_loss import DiouLoss
+from .ops import BBoxAssigner, LibraBBoxAssigner
 
 __all__ = ['MaskRCNN']
 
@@ -46,11 +48,11 @@ class MaskRCNN(object):
     def __init__(
             self,
             backbone,
-            num_classes=81,
+            input_channel=3,
             mode='train',
+            num_classes=81,
             with_fpn=False,
             fpn=None,
-            num_chan=256,
             min_level=2,
             max_level=6,
             spatial_scale=[1. / 32., 1. / 16., 1. / 8., 1. / 4.],
@@ -69,6 +71,9 @@ class MaskRCNN(object):
             test_pre_nms_top_n=6000,
             test_post_nms_top_n=1000,
             test_nms_thresh=0.7,
+            rpn_cls_loss='SigmoidCrossEntropy',
+            rpn_focal_loss_alpha=0.25,
+            rpn_focal_loss_gamma=2,
             #roi_extractor
             roi_extractor=None,
             #bbox_head
@@ -76,6 +81,10 @@ class MaskRCNN(object):
             keep_top_k=100,
             nms_threshold=0.5,
             score_threshold=0.05,
+            rcnn_nms='MultiClassNMS',
+            softnms_sigma=0.5,
+            post_threshold=.05,
+
             #MaskHead
             mask_head=None,
             num_convs=0,
@@ -88,7 +97,12 @@ class MaskRCNN(object):
             bg_thresh_lo=0.,
             bbox_reg_weights=[0.1, 0.1, 0.2, 0.2],
             fixed_input_shape=None,
-            input_channel=3):
+            rcnn_bbox_loss='SmoothL1Loss',
+            diouloss_weight=10.0,
+            diouloss_is_cls_agnostic=False,
+            diouloss_use_complete_iou_loss=True,
+            bbox_assigner='BBoxAssigner',
+            fpn_num_channels=256):
         super(MaskRCNN, self).__init__()
         self.backbone = backbone
         self.mode = mode
@@ -98,11 +112,13 @@ class MaskRCNN(object):
                 fpn.min_level = 2
                 fpn.max_level = 6
             else:
-                fpn = FPN(num_chan=num_chan,
+                fpn = FPN(num_chan=fpn_num_channels,
                           min_level=min_level,
                           max_level=max_level,
                           spatial_scale=spatial_scale)
         self.fpn = fpn
+        if self.fpn is not None:
+            self.fpn.num_chan = fpn_num_channels
         self.num_classes = num_classes
         if rpn_head is None:
             if self.fpn is None:
@@ -118,7 +134,10 @@ class MaskRCNN(object):
                     train_nms_thresh=train_nms_thresh,
                     test_pre_nms_top_n=test_pre_nms_top_n,
                     test_post_nms_top_n=test_post_nms_top_n,
-                    test_nms_thresh=test_nms_thresh)
+                    test_nms_thresh=test_nms_thresh,
+                    rpn_cls_loss=rpn_cls_loss,
+                    rpn_focal_loss_alpha=rpn_focal_loss_alpha,
+                    rpn_focal_loss_gamma=rpn_focal_loss_gamma)
             else:
                 rpn_head = FPNRPNHead(
                     anchor_start_size=anchor_sizes[0],
@@ -135,7 +154,10 @@ class MaskRCNN(object):
                     train_nms_thresh=train_nms_thresh,
                     test_pre_nms_top_n=test_pre_nms_top_n,
                     test_post_nms_top_n=test_post_nms_top_n,
-                    test_nms_thresh=test_nms_thresh)
+                    test_nms_thresh=test_nms_thresh,
+                    rpn_cls_loss=rpn_cls_loss,
+                    rpn_focal_loss_alpha=rpn_focal_loss_alpha,
+                    rpn_focal_loss_gamma=rpn_focal_loss_gamma)
         self.rpn_head = rpn_head
         if roi_extractor is None:
             if self.fpn is None:
@@ -150,7 +172,8 @@ class MaskRCNN(object):
                 head = ResNetC5(
                     layers=self.backbone.layers,
                     norm_type=self.backbone.norm_type,
-                    freeze_norm=self.backbone.freeze_norm)
+                    freeze_norm=self.backbone.freeze_norm,
+                    variant=self.backbone.variant)
             else:
                 head = TwoFCHead()
             bbox_head = BBoxHead(
@@ -158,7 +181,15 @@ class MaskRCNN(object):
                 keep_top_k=keep_top_k,
                 nms_threshold=nms_threshold,
                 score_threshold=score_threshold,
-                num_classes=num_classes)
+                rcnn_nms=rcnn_nms,
+                softnms_sigma=softnms_sigma,
+                post_threshold=post_threshold,
+                num_classes=num_classes,
+                rcnn_bbox_loss=rcnn_bbox_loss,
+                diouloss_weight=diouloss_weight,
+                diouloss_is_cls_agnostic=diouloss_is_cls_agnostic,
+                diouloss_use_complete_iou_loss=diouloss_use_complete_iou_loss)
+
         self.bbox_head = bbox_head
         if mask_head is None:
             mask_head = MaskHead(
@@ -174,14 +205,40 @@ class MaskRCNN(object):
         self.bbox_reg_weights = bbox_reg_weights
         self.rpn_only = rpn_only
         self.fixed_input_shape = fixed_input_shape
+        if bbox_assigner == 'BBoxAssigner':
+            self.bbox_assigner = BBoxAssigner(
+                batch_size_per_im=batch_size_per_im,
+                fg_fraction=fg_fraction,
+                fg_thresh=fg_thresh,
+                bg_thresh_hi=bg_thresh_hi,
+                bg_thresh_lo=bg_thresh_lo,
+                bbox_reg_weights=bbox_reg_weights,
+                num_classes=num_classes,
+                shuffle_before_sample=self.rpn_head.use_random)
+        elif bbox_assigner == 'LibraBBoxAssigner':
+            self.bbox_assigner = LibraBBoxAssigner(
+                batch_size_per_im=batch_size_per_im,
+                fg_fraction=fg_fraction,
+                fg_thresh=fg_thresh,
+                bg_thresh_hi=bg_thresh_hi,
+                bg_thresh_lo=bg_thresh_lo,
+                bbox_reg_weights=bbox_reg_weights,
+                num_classes=num_classes,
+                shuffle_before_sample=self.rpn_head.use_random)
         self.input_channel = input_channel
 
     def build_net(self, inputs):
         im = inputs['image']
         im_info = inputs['im_info']
+        if self.mode == 'train':
+            gt_bbox = inputs['gt_box']
+            is_crowd = inputs['is_crowd']
+        else:
+            im_shape = inputs['im_shape']
 
         # backbone
         body_feats = self.backbone(im)
+        body_feat_names = list(body_feats.keys())
 
         # FPN
         spatial_scale = None
@@ -192,28 +249,19 @@ class MaskRCNN(object):
         rois = self.rpn_head.get_proposals(body_feats, im_info, mode=self.mode)
 
         if self.mode == 'train':
-            rpn_loss = self.rpn_head.get_loss(im_info, inputs['gt_box'],
-                                              inputs['is_crowd'])
-            outputs = fluid.layers.generate_proposal_labels(
+            rpn_loss = self.rpn_head.get_loss(im_info, gt_bbox, is_crowd)
+            outputs = self.bbox_assigner(
                 rpn_rois=rois,
                 gt_classes=inputs['gt_label'],
                 is_crowd=inputs['is_crowd'],
                 gt_boxes=inputs['gt_box'],
-                im_info=inputs['im_info'],
-                batch_size_per_im=self.batch_size_per_im,
-                fg_fraction=self.fg_fraction,
-                fg_thresh=self.fg_thresh,
-                bg_thresh_hi=self.bg_thresh_hi,
-                bg_thresh_lo=self.bg_thresh_lo,
-                bbox_reg_weights=self.bbox_reg_weights,
-                class_nums=self.num_classes,
-                use_random=self.rpn_head.use_random)
+                im_info=inputs['im_info'])
 
             rois = outputs[0]
             labels_int32 = outputs[1]
 
             if self.fpn is None:
-                last_feat = body_feats[list(body_feats.keys())[-1]]
+                last_feat = body_feats[body_feat_names[-1]]
                 roi_feat = self.roi_extractor(last_feat, rois)
             else:
                 roi_feat = self.roi_extractor(body_feats, rois, spatial_scale)
@@ -254,8 +302,7 @@ class MaskRCNN(object):
                 return {'proposal': rois}
             mask_name = 'mask_pred'
             mask_pred, bbox_pred = self._eval(body_feats, mask_name, rois,
-                                              im_info, inputs['im_shape'],
-                                              spatial_scale)
+                                              im_info, im_shape, spatial_scale)
             return OrderedDict(zip(['bbox', 'mask'], [bbox_pred, mask_pred]))
 
     def _eval(self,