fix save_model bug and remove inputs[mode] in rcnn (#2031)

dygraph/ppdet/modeling/architectures/cascade_rcnn.py

@@ -99,6 +99,7 @@ class CascadeRCNN(BaseArch):
                 self.inputs,
                 self.rpn_head_out,
                 self.anchor_out,
+                self.training,
                 i,
                 rois,
                 bbox_head_out,
@@ -110,7 +111,7 @@ class CascadeRCNN(BaseArch):
                                                          spatial_scale, i)
             self.bbox_head_list.append(bbox_head_out)
 
-        if self.inputs['mode'] == 'infer':
+        if not self.training:
             bbox_pred, bboxes = self.bbox_head.get_cascade_prediction(
                 self.bbox_head_list, rois_list)
             self.bboxes = self.bbox_post_process(bbox_pred, bboxes,
@@ -120,7 +121,7 @@ class CascadeRCNN(BaseArch):
         if self.with_mask:
             rois = rois_list[-1]
             rois_has_mask_int32 = None
-            if self.inputs['mode'] == 'train':
+            if self.training:
                 bbox_targets = self.proposal.get_targets()[-1]
                 self.bboxes, rois_has_mask_int32 = self.mask(self.inputs, rois,
                                                              bbox_targets)

dygraph/ppdet/modeling/architectures/faster_rcnn.py

@@ -57,12 +57,13 @@ class FasterRCNN(BaseArch):
 
         # Proposal RoI
         # compute targets here when training
-        rois = self.proposal(self.inputs, self.rpn_head_out, self.anchor_out)
+        rois = self.proposal(self.inputs, self.rpn_head_out, self.anchor_out,
+                             self.training)
         # BBox Head
         bbox_feat, self.bbox_head_out, self.bbox_head_feat_func = self.bbox_head(
             body_feats, rois, spatial_scale)
 
-        if self.inputs['mode'] == 'infer':
+        if not self.training:
             bbox_pred, bboxes = self.bbox_head.get_prediction(
                 self.bbox_head_out, rois)
             # Refine bbox by the output from bbox_head at test stage

dygraph/ppdet/modeling/architectures/mask_rcnn.py

@@ -85,13 +85,14 @@ class MaskRCNN(BaseArch):
 
         # Proposal RoI 
         # compute targets here when training
-        rois = self.proposal(self.inputs, self.rpn_head_out, self.anchor_out)
+        rois = self.proposal(self.inputs, self.rpn_head_out, self.anchor_out,
+                             self.training)
         # BBox Head
         bbox_feat, self.bbox_head_out, bbox_head_feat_func = self.bbox_head(
             body_feats, rois, spatial_scale)
 
         rois_has_mask_int32 = None
-        if self.inputs['mode'] == 'infer':
+        if not self.training:
             bbox_pred, bboxes = self.bbox_head.get_prediction(
                 self.bbox_head_out, rois)
             # Refine bbox by the output from bbox_head at test stage

dygraph/ppdet/modeling/architectures/meta_arch.py

@@ -15,18 +15,14 @@ class BaseArch(nn.Layer):
     def __init__(self):
         super(BaseArch, self).__init__()
 
-    def forward(self, inputs, mode='infer'):
+    def forward(self, inputs):
         self.inputs = inputs
-        self.inputs['mode'] = mode
         self.model_arch()
 
-        if mode == 'train':
+        if self.training:
             out = self.get_loss()
-        elif mode == 'infer':
-            out = self.get_pred()
         else:
-            out = None
-            raise "Now, only support train and infer mode!"
+            out = self.get_pred()
         return out
 
     def build_inputs(self, data, input_def):

dygraph/ppdet/modeling/bbox.py

@@ -79,7 +79,7 @@ class Proposal(object):
         self.proposal_generator = proposal_generator
         self.proposal_target_generator = proposal_target_generator
 
-    def generate_proposal(self, inputs, rpn_head_out, anchor_out):
+    def generate_proposal(self, inputs, rpn_head_out, anchor_out, is_train):
         # TODO: delete im_info 
         try:
             im_shape = inputs['im_info']
@@ -97,7 +97,7 @@ class Proposal(object):
                 anchors=anchor,
                 variances=var,
                 im_shape=im_shape,
-                mode=inputs['mode'])
+                is_train=is_train)
             if len(rpn_head_out) == 1:
                 return rpn_rois, rpn_rois_num
             rpn_rois_list.append(rpn_rois)
@@ -164,13 +164,14 @@ class Proposal(object):
                  inputs,
                  rpn_head_out,
                  anchor_out,
+                 is_train=False,
                  stage=0,
                  proposal_out=None,
                  bbox_head_out=None,
                  max_overlap=None):
         if stage == 0:
             roi, rois_num = self.generate_proposal(inputs, rpn_head_out,
-                                                   anchor_out)
+                                                   anchor_out, is_train)
             self.targets_list = []
             self.max_overlap = None
 
@@ -178,7 +179,7 @@ class Proposal(object):
             bbox_delta = bbox_head_out[1]
             roi = self.refine_bbox(proposal_out[0], bbox_delta, stage)
             rois_num = proposal_out[1]
-        if inputs['mode'] == 'train':
+        if is_train:
             roi, rois_num, targets, self.max_overlap = self.generate_proposal_target(
                 inputs, roi, rois_num, stage, self.max_overlap)
             self.targets_list.append(targets)

dygraph/ppdet/modeling/heads/mask_head.py

@@ -83,14 +83,13 @@ class MaskFeat(Layer):
                 mask_index,
                 spatial_scale,
                 stage=0,
-                bbox_head_feat_func=None,
-                mode='train'):
+                bbox_head_feat_func=None):
         if self.share_bbox_feat and mask_index is not None:
             rois_feat = paddle.gather(bbox_feat, mask_index)
         else:
             rois_feat = self.mask_roi_extractor(body_feats, bboxes,
                                                 spatial_scale)
-        if self.share_bbox_feat and bbox_head_feat_func is not None and mode == 'infer':
+        if self.share_bbox_feat and bbox_head_feat_func is not None and not self.training:
             rois_feat = bbox_head_feat_func(rois_feat)
 
         # upsample 
@@ -136,14 +135,8 @@ class MaskHead(Layer):
                       spatial_scale,
                       stage=0):
         # feat
-        mask_feat = self.mask_feat(
-            body_feats,
-            bboxes,
-            bbox_feat,
-            mask_index,
-            spatial_scale,
-            stage,
-            mode='train')
+        mask_feat = self.mask_feat(body_feats, bboxes, bbox_feat, mask_index,
+                                   spatial_scale, stage)
         # logits
         mask_head_out = self.mask_fcn_logits[stage](mask_feat)
         return mask_head_out
@@ -174,15 +167,9 @@ class MaskHead(Layer):
             scale_factor_list = paddle.reshape(scale_factor_list, shape=[-1, 1])
             scaled_bbox = paddle.multiply(bbox[:, 2:], scale_factor_list)
             scaled_bboxes = (scaled_bbox, bbox_num)
-            mask_feat = self.mask_feat(
-                body_feats,
-                scaled_bboxes,
-                bbox_feat,
-                mask_index,
-                spatial_scale,
-                stage,
-                bbox_head_feat_func,
-                mode='infer')
+            mask_feat = self.mask_feat(body_feats, scaled_bboxes, bbox_feat,
+                                       mask_index, spatial_scale, stage,
+                                       bbox_head_feat_func)
             mask_logit = self.mask_fcn_logits[stage](mask_feat)
             mask_head_out = F.sigmoid(mask_logit)
         return mask_head_out
@@ -196,7 +183,7 @@ class MaskHead(Layer):
                 spatial_scale,
                 bbox_head_feat_func=None,
                 stage=0):
-        if inputs['mode'] == 'train':
+        if self.training:
             mask_head_out = self.forward_train(body_feats, bboxes, bbox_feat,
                                                mask_index, spatial_scale, stage)
         else:

dygraph/ppdet/modeling/layers.py

@@ -189,9 +189,9 @@ class ProposalGenerator(object):
                  anchors,
                  variances,
                  im_shape,
-                 mode='train'):
-        pre_nms_top_n = self.train_pre_nms_top_n if mode == 'train' else self.infer_pre_nms_top_n
-        post_nms_top_n = self.train_post_nms_top_n if mode == 'train' else self.infer_post_nms_top_n
+                 is_train=False):
+        pre_nms_top_n = self.train_pre_nms_top_n if is_train else self.infer_pre_nms_top_n
+        post_nms_top_n = self.train_post_nms_top_n if is_train else self.infer_post_nms_top_n
         # TODO delete im_info
         if im_shape.shape[1] > 2:
             import paddle.fluid as fluid

dygraph/ppdet/utils/checkpoint.py

@@ -159,6 +159,8 @@ def save_model(model, optimizer, save_dir, save_name, last_epoch):
         save_name (str): the path to be saved.
         last_epoch (int): the epoch index.
     """
+    if paddle.distributed.get_rank() != 0:
+        return
     if not os.path.exists(save_dir):
         os.makedirs(save_dir)
     save_path = os.path.join(save_dir, save_name)

dygraph/tools/eval.py

@@ -81,7 +81,7 @@ def run(FLAGS, cfg, place):
     for iter_id, data in enumerate(eval_loader):
         # forward
         model.eval()
-        outs = model(data, mode='infer')
+        outs = model(data)
         for key in extra_key:
             outs[key] = data[key]
         for key, value in outs.items():

dygraph/tools/infer.py

@@ -153,7 +153,7 @@ def run(FLAGS, cfg, place):
     for iter_id, data in enumerate(test_loader):
         # forward
         model.eval()
-        outs = model(data, mode='infer')
+        outs = model(data)
         for key in extra_key:
             outs[key] = data[key]
         for key, value in outs.items():

dygraph/tools/train.py

@@ -164,7 +164,7 @@ def run(FLAGS, cfg, place):
             data_time.update(time.time() - end_time)
             # Model Forward
             model.train()
-            outputs = model(data, mode='train')
+            outputs = model(data)
             loss = outputs['loss']
             # Model Backward
             loss.backward()