diff --git a/ppdet/modeling/roi_heads/bbox_head.py b/ppdet/modeling/roi_heads/bbox_head.py index ce2fb63d756b378cba026a41ce329cc310cbcffc..16a7db529dfd8f0740ceae04210a86a9b3e6ac35 100644 --- a/ppdet/modeling/roi_heads/bbox_head.py +++ b/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: