From 636b8c4794202470e310c8af07c9a632b7e90671 Mon Sep 17 00:00:00 2001 From: wangguanzhong Date: Thu, 9 Jun 2022 18:06:56 +0800 Subject: [PATCH] fix mask rcnn in eval when num_classes is 1 (#6168) --- ppdet/modeling/heads/mask_head.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ppdet/modeling/heads/mask_head.py b/ppdet/modeling/heads/mask_head.py index 08b7d3d26..939debbaa 100644 --- a/ppdet/modeling/heads/mask_head.py +++ b/ppdet/modeling/heads/mask_head.py @@ -222,7 +222,7 @@ class MaskHead(nn.Layer): mask_feat = self.head(rois_feat) mask_logit = self.mask_fcn_logits(mask_feat) if self.num_classes == 1: - mask_out = F.sigmoid(mask_logit) + mask_out = F.sigmoid(mask_logit)[:, 0, :, :] else: num_masks = paddle.shape(mask_logit)[0] index = paddle.arange(num_masks).cast('int32') -- GitLab