未验证 提交 7042801c 编写于 作者: G Guanghua Yu 提交者: GitHub

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

上级 28510251
......@@ -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)
......
......@@ -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
......
......@@ -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
......
......@@ -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):
......
......@@ -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)
......
......@@ -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:
......
......@@ -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
......
......@@ -150,7 +150,7 @@ def load_pretrain_weight(model,
def save_model(model, optimizer, save_dir, save_name, last_epoch):
"""
save model into disk.
Args:
model (paddle.nn.Layer): the Layer instalce to save parameters.
optimizer (paddle.optimizer.Optimizer): the Optimizer instance to
......@@ -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)
......
......@@ -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():
......
......@@ -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():
......
......@@ -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()
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册