Use fp16 for RCNN bbox head (#3475)

ppdet/modeling/roi_heads/bbox_head.py

@@ -27,6 +27,7 @@ from paddle.fluid.initializer import MSRA
 from ppdet.modeling.ops import MultiClassNMS
 from ppdet.modeling.ops import ConvNorm
 from ppdet.core.workspace import register, serializable
+from ppdet.experimental import mixed_precision_global_state
 
 __all__ = ['BBoxHead', 'TwoFCHead', 'XConvNormHead']
 
@@ -120,6 +121,12 @@ class TwoFCHead(object):
 
     def __call__(self, roi_feat):
         fan = roi_feat.shape[1] * roi_feat.shape[2] * roi_feat.shape[3]
+
+        mixed_precision_enabled = mixed_precision_global_state() is not None
+
+        if mixed_precision_enabled:
+            roi_feat = fluid.layers.cast(roi_feat, 'float16')
+
         fc6 = fluid.layers.fc(input=roi_feat,
                               size=self.mlp_dim,
                               act='relu',
@@ -141,6 +148,10 @@ class TwoFCHead(object):
                                         name='fc7_b',
                                         learning_rate=2.,
                                         regularizer=L2Decay(0.)))
+
+        if mixed_precision_enabled:
+            head_feat = fluid.layers.cast(head_feat, 'float32')
+
         return head_feat
 
 
@@ -280,7 +291,7 @@ class BBoxHead(object):
                 number of input images, each element consists of im_height,
                 im_width, im_scale.
             im_shape (Variable): Actual shape of original image with shape
-                [B, 3]. B is the number of images, each element consists of 
+                [B, 3]. B is the number of images, each element consists of
                 original_height, original_width, 1
 
         Returns: