def forward_train(self,
img,
img_meta,
gt_bboxes,
gt_labels,
gt_bboxes_ignore=None,
gt_masks=None,
proposals=None):"""
Args:
img (Tensor): of shape (N, C, H, W) encoding input images.
Typically these should be mean centered and std scaled.
img_meta (list[dict]): list of image info dict where each dict has:'img_shape','scale_factor','flip',and may also contain
'filename','ori_shape','pad_shape',and'img_norm_cfg'.
For details on the values of these keys see
`mmdet/datasets/pipelines/formatting.py:Collect`.
gt_bboxes (list[Tensor]): each item are the truth boxes for each
image in [tl_x, tl_y, br_x, br_y] format.
gt_labels (list[Tensor]):classindices corresponding to each box
gt_bboxes_ignore (None | list[Tensor]): specify which bounding
boxes can be ignored when computing the loss.
gt_masks (None | Tensor):true segmentation masks for each box
used if the architecture supports a segmentation task.
proposals : override rpn proposals with custom proposals. Use when
`with_rpn` is False.
Returns:
dict[str, Tensor]: a dictionary of loss components
"""
x = self.extract_feat(img) # img ->|backbone+neck|-> base_feature
losses =dict()# RPN forward and loss,RPN前向传播部分,生成proposalif self.with_rpn:
rpn_outs = self.rpn_head(x) # rpn_outs =((rpn_cls_score1, rpn_cls_score2,..., rpn_cls_score5),
# (rpn_bbox_pred1, rpn_bbox_pred2,..., rpn_bbox_pred5))# shape: 图片数量*fpn层数*channel*H*W, 图片数量*fpn层数*channel*H*W
rpn_loss_inputs = rpn_outs +(gt_bboxes, img_meta,
self.train_cfg.rpn) # 用于计算rpn loss的输入
rpn_losses = self.rpn_head.loss(*rpn_loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore) # 计算loss
losses.update(rpn_losses) # 更新loss
proposal_cfg = self.train_cfg.get('rpn_proposal',
self.test_cfg.rpn)
proposal_inputs = rpn_outs +(img_meta, proposal_cfg)
proposal_list = self.rpn_head.get_bboxes(*proposal_inputs) # rpn输出的proposal
else:
proposal_list = proposals
# assign gts and sample proposals,将proposal与ground truth匹配if self.with_bbox or self.with_mask:
bbox_assigner =build_assigner(self.train_cfg.rcnn.assigner)
bbox_sampler =build_sampler(
self.train_cfg.rcnn.sampler, context=self)
num_imgs = img.size(0)if gt_bboxes_ignore is None:
gt_bboxes_ignore =[None for _ in range(num_imgs)]
sampling_results =[]for i in range(num_imgs):
assign_result = bbox_assigner.assign(proposal_list[i],
gt_bboxes[i],
gt_bboxes_ignore[i],
gt_labels[i])
sampling_result = bbox_sampler.sample(
assign_result,
proposal_list[i],
gt_bboxes[i],
gt_labels[i],
feats=[lvl_feat[i][None]for lvl_feat in x])
sampling_results.append(sampling_result)# bbox head forward and loss,对proposal特征进行池化操作,经过分类分支和回归分支后分别计算分类误差和回归误差if self.with_bbox:
rois =bbox2roi([res.bboxes for res in sampling_results]) # 为每个proposal添加图片索引号
# TODO: a more flexible way to decide which feature maps to use
bbox_feats = self.bbox_roi_extractor(
x[:self.bbox_roi_extractor.num_inputs], rois) # rois池化后的特征
if self.with_shared_head:
bbox_feats = self.shared_head(bbox_feats)
cls_score, bbox_pred = self.bbox_head(bbox_feats)
bbox_targets = self.bbox_head.get_target(sampling_results,
gt_bboxes, gt_labels,
self.train_cfg.rcnn)
loss_bbox = self.bbox_head.loss(cls_score, bbox_pred,*bbox_targets)
losses.update(loss_bbox)# mask head forward and lossif self.with_mask:ifnot self.share_roi_extractor:
pos_rois =bbox2roi([res.pos_bboxes for res in sampling_results])
mask_feats = self.mask_roi_extractor(
x[:self.mask_roi_extractor.num_inputs], pos_rois)if self.with_shared_head:
mask_feats = self.shared_head(mask_feats)else:
pos_inds =[]
device = bbox_feats.device
for res in sampling_results:
pos_inds.append(
torch.ones(
res.pos_bboxes.shape[0],
device=device,
dtype=torch.uint8))
pos_inds.append(
torch.zeros(
res.neg_bboxes.shape[0],
device=device,
dtype=torch.uint8))
pos_inds = torch.cat(pos_inds)
mask_feats = bbox_feats[pos_inds]if mask_feats.shape[0]>0:
mask_pred = self.mask_head(mask_feats)
mask_targets = self.mask_head.get_target(
sampling_results, gt_masks, self.train_cfg.rcnn)
pos_labels = torch.cat([res.pos_gt_labels for res in sampling_results])
loss_mask = self.mask_head.loss(mask_pred, mask_targets,
pos_labels)
losses.update(loss_mask)return losses