ppyoloe_head.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385
  1. # Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import paddle
  15. import paddle.nn as nn
  16. import paddle.nn.functional as F
  17. from paddlex.ppdet.core.workspace import register
  18. from ..bbox_utils import batch_distance2bbox
  19. from ..losses import GIoULoss
  20. from ..initializer import bias_init_with_prob, constant_, normal_
  21. from ..assigners.utils import generate_anchors_for_grid_cell
  22. from paddlex.ppdet.modeling.backbones.cspresnet import ConvBNLayer
  23. from paddlex.ppdet.modeling.ops import get_static_shape, get_act_fn
  24. from paddlex.ppdet.modeling.layers import MultiClassNMS
  25. __all__ = ['PPYOLOEHead']
  26. class ESEAttn(nn.Layer):
  27. def __init__(self, feat_channels, act='swish'):
  28. super(ESEAttn, self).__init__()
  29. self.fc = nn.Conv2D(feat_channels, feat_channels, 1)
  30. self.conv = ConvBNLayer(feat_channels, feat_channels, 1, act=act)
  31. self._init_weights()
  32. def _init_weights(self):
  33. normal_(self.fc.weight, std=0.001)
  34. def forward(self, feat, avg_feat):
  35. weight = F.sigmoid(self.fc(avg_feat))
  36. return self.conv(feat * weight)
  37. @register
  38. class PPYOLOEHead(nn.Layer):
  39. __shared__ = ['num_classes', 'eval_size', 'trt', 'exclude_nms']
  40. __inject__ = ['static_assigner', 'assigner', 'nms']
  41. def __init__(self,
  42. in_channels=[1024, 512, 256],
  43. num_classes=80,
  44. act='swish',
  45. fpn_strides=(32, 16, 8),
  46. grid_cell_scale=5.0,
  47. grid_cell_offset=0.5,
  48. reg_max=16,
  49. static_assigner_epoch=4,
  50. use_varifocal_loss=True,
  51. static_assigner='ATSSAssigner',
  52. assigner='TaskAlignedAssigner',
  53. nms='MultiClassNMS',
  54. eval_size=None,
  55. loss_weight={
  56. 'class': 1.0,
  57. 'iou': 2.5,
  58. 'dfl': 0.5,
  59. },
  60. trt=False,
  61. exclude_nms=False):
  62. super(PPYOLOEHead, self).__init__()
  63. assert len(in_channels) > 0, "len(in_channels) should > 0"
  64. self.in_channels = in_channels
  65. self.num_classes = num_classes
  66. self.fpn_strides = fpn_strides
  67. self.grid_cell_scale = grid_cell_scale
  68. self.grid_cell_offset = grid_cell_offset
  69. self.reg_max = reg_max
  70. self.iou_loss = GIoULoss()
  71. self.loss_weight = loss_weight
  72. self.use_varifocal_loss = use_varifocal_loss
  73. self.eval_size = eval_size
  74. self.static_assigner_epoch = static_assigner_epoch
  75. self.static_assigner = static_assigner
  76. self.assigner = assigner
  77. self.nms = nms
  78. if isinstance(self.nms, MultiClassNMS) and trt:
  79. self.nms.trt = trt
  80. self.exclude_nms = exclude_nms
  81. # stem
  82. self.stem_cls = nn.LayerList()
  83. self.stem_reg = nn.LayerList()
  84. act = get_act_fn(
  85. act, trt=trt) if act is None or isinstance(act,
  86. (str, dict)) else act
  87. for in_c in self.in_channels:
  88. self.stem_cls.append(ESEAttn(in_c, act=act))
  89. self.stem_reg.append(ESEAttn(in_c, act=act))
  90. # pred head
  91. self.pred_cls = nn.LayerList()
  92. self.pred_reg = nn.LayerList()
  93. for in_c in self.in_channels:
  94. self.pred_cls.append(
  95. nn.Conv2D(
  96. in_c, self.num_classes, 3, padding=1))
  97. self.pred_reg.append(
  98. nn.Conv2D(
  99. in_c, 4 * (self.reg_max + 1), 3, padding=1))
  100. # projection conv
  101. self.proj_conv = nn.Conv2D(self.reg_max + 1, 1, 1, bias_attr=False)
  102. self.proj_conv.skip_quant = True
  103. self._init_weights()
  104. @classmethod
  105. def from_config(cls, cfg, input_shape):
  106. return {'in_channels': [i.channels for i in input_shape], }
  107. def _init_weights(self):
  108. bias_cls = bias_init_with_prob(0.01)
  109. for cls_, reg_ in zip(self.pred_cls, self.pred_reg):
  110. constant_(cls_.weight)
  111. constant_(cls_.bias, bias_cls)
  112. constant_(reg_.weight)
  113. constant_(reg_.bias, 1.0)
  114. self.proj = paddle.linspace(0, self.reg_max, self.reg_max + 1)
  115. self.proj_conv.weight.set_value(
  116. self.proj.reshape([1, self.reg_max + 1, 1, 1]))
  117. self.proj_conv.weight.stop_gradient = True
  118. if self.eval_size:
  119. anchor_points, stride_tensor = self._generate_anchors()
  120. self.anchor_points = anchor_points
  121. self.stride_tensor = stride_tensor
  122. def forward_train(self, feats, targets):
  123. anchors, anchor_points, num_anchors_list, stride_tensor = \
  124. generate_anchors_for_grid_cell(
  125. feats, self.fpn_strides, self.grid_cell_scale,
  126. self.grid_cell_offset)
  127. cls_score_list, reg_distri_list = [], []
  128. for i, feat in enumerate(feats):
  129. avg_feat = F.adaptive_avg_pool2d(feat, (1, 1))
  130. cls_logit = self.pred_cls[i](self.stem_cls[i](feat, avg_feat) +
  131. feat)
  132. reg_distri = self.pred_reg[i](self.stem_reg[i](feat, avg_feat))
  133. # cls and reg
  134. cls_score = F.sigmoid(cls_logit)
  135. cls_score_list.append(cls_score.flatten(2).transpose([0, 2, 1]))
  136. reg_distri_list.append(reg_distri.flatten(2).transpose([0, 2, 1]))
  137. cls_score_list = paddle.concat(cls_score_list, axis=1)
  138. reg_distri_list = paddle.concat(reg_distri_list, axis=1)
  139. return self.get_loss([
  140. cls_score_list, reg_distri_list, anchors, anchor_points,
  141. num_anchors_list, stride_tensor
  142. ], targets)
  143. def _generate_anchors(self, feats=None):
  144. # just use in eval time
  145. anchor_points = []
  146. stride_tensor = []
  147. for i, stride in enumerate(self.fpn_strides):
  148. if feats is not None:
  149. _, _, h, w = feats[i].shape
  150. else:
  151. h = int(self.eval_size[0] / stride)
  152. w = int(self.eval_size[1] / stride)
  153. shift_x = paddle.arange(end=w) + self.grid_cell_offset
  154. shift_y = paddle.arange(end=h) + self.grid_cell_offset
  155. shift_y, shift_x = paddle.meshgrid(shift_y, shift_x)
  156. anchor_point = paddle.cast(
  157. paddle.stack(
  158. [shift_x, shift_y], axis=-1), dtype='float32')
  159. anchor_points.append(anchor_point.reshape([-1, 2]))
  160. stride_tensor.append(
  161. paddle.full(
  162. [h * w, 1], stride, dtype='float32'))
  163. anchor_points = paddle.concat(anchor_points)
  164. stride_tensor = paddle.concat(stride_tensor)
  165. return anchor_points, stride_tensor
  166. def forward_eval(self, feats):
  167. if self.eval_size:
  168. anchor_points, stride_tensor = self.anchor_points, self.stride_tensor
  169. else:
  170. anchor_points, stride_tensor = self._generate_anchors(feats)
  171. cls_score_list, reg_dist_list = [], []
  172. for i, feat in enumerate(feats):
  173. b, _, h, w = feat.shape
  174. l = h * w
  175. avg_feat = F.adaptive_avg_pool2d(feat, (1, 1))
  176. cls_logit = self.pred_cls[i](self.stem_cls[i](feat, avg_feat) +
  177. feat)
  178. reg_dist = self.pred_reg[i](self.stem_reg[i](feat, avg_feat))
  179. reg_dist = reg_dist.reshape(
  180. [-1, 4, self.reg_max + 1, l]).transpose([0, 2, 1, 3])
  181. reg_dist = self.proj_conv(F.softmax(reg_dist, axis=1))
  182. # cls and reg
  183. cls_score = F.sigmoid(cls_logit)
  184. cls_score_list.append(cls_score.reshape([b, self.num_classes, l]))
  185. reg_dist_list.append(reg_dist.reshape([b, 4, l]))
  186. cls_score_list = paddle.concat(cls_score_list, axis=-1)
  187. reg_dist_list = paddle.concat(reg_dist_list, axis=-1)
  188. return cls_score_list, reg_dist_list, anchor_points, stride_tensor
  189. def forward(self, feats, targets=None):
  190. assert len(feats) == len(self.fpn_strides), \
  191. "The size of feats is not equal to size of fpn_strides"
  192. if self.training:
  193. return self.forward_train(feats, targets)
  194. else:
  195. return self.forward_eval(feats)
  196. @staticmethod
  197. def _focal_loss(score, label, alpha=0.25, gamma=2.0):
  198. weight = (score - label).pow(gamma)
  199. if alpha > 0:
  200. alpha_t = alpha * label + (1 - alpha) * (1 - label)
  201. weight *= alpha_t
  202. loss = F.binary_cross_entropy(
  203. score, label, weight=weight, reduction='sum')
  204. return loss
  205. @staticmethod
  206. def _varifocal_loss(pred_score, gt_score, label, alpha=0.75, gamma=2.0):
  207. weight = alpha * pred_score.pow(gamma) * (1 - label) + gt_score * label
  208. loss = F.binary_cross_entropy(
  209. pred_score, gt_score, weight=weight, reduction='sum')
  210. return loss
  211. def _bbox_decode(self, anchor_points, pred_dist):
  212. b, l, _ = get_static_shape(pred_dist)
  213. pred_dist = F.softmax(pred_dist.reshape([b, l, 4, self.reg_max + 1
  214. ])).matmul(self.proj)
  215. return batch_distance2bbox(anchor_points, pred_dist)
  216. def _bbox2distance(self, points, bbox):
  217. x1y1, x2y2 = paddle.split(bbox, 2, -1)
  218. lt = points - x1y1
  219. rb = x2y2 - points
  220. return paddle.concat([lt, rb], -1).clip(0, self.reg_max - 0.01)
  221. def _df_loss(self, pred_dist, target):
  222. target_left = paddle.cast(target, 'int64')
  223. target_right = target_left + 1
  224. weight_left = target_right.astype('float32') - target
  225. weight_right = 1 - weight_left
  226. loss_left = F.cross_entropy(
  227. pred_dist, target_left, reduction='none') * weight_left
  228. loss_right = F.cross_entropy(
  229. pred_dist, target_right, reduction='none') * weight_right
  230. return (loss_left + loss_right).mean(-1, keepdim=True)
  231. def _bbox_loss(self, pred_dist, pred_bboxes, anchor_points,
  232. assigned_labels, assigned_bboxes, assigned_scores,
  233. assigned_scores_sum):
  234. # select positive samples mask
  235. mask_positive = (assigned_labels != self.num_classes)
  236. num_pos = mask_positive.sum()
  237. # pos/neg loss
  238. if num_pos > 0:
  239. # l1 + iou
  240. bbox_mask = mask_positive.unsqueeze(-1).tile([1, 1, 4])
  241. pred_bboxes_pos = paddle.masked_select(pred_bboxes,
  242. bbox_mask).reshape([-1, 4])
  243. assigned_bboxes_pos = paddle.masked_select(
  244. assigned_bboxes, bbox_mask).reshape([-1, 4])
  245. bbox_weight = paddle.masked_select(
  246. assigned_scores.sum(-1), mask_positive).unsqueeze(-1)
  247. loss_l1 = F.l1_loss(pred_bboxes_pos, assigned_bboxes_pos)
  248. loss_iou = self.iou_loss(pred_bboxes_pos,
  249. assigned_bboxes_pos) * bbox_weight
  250. loss_iou = loss_iou.sum() / assigned_scores_sum
  251. dist_mask = mask_positive.unsqueeze(-1).tile(
  252. [1, 1, (self.reg_max + 1) * 4])
  253. pred_dist_pos = paddle.masked_select(
  254. pred_dist, dist_mask).reshape([-1, 4, self.reg_max + 1])
  255. assigned_ltrb = self._bbox2distance(anchor_points, assigned_bboxes)
  256. assigned_ltrb_pos = paddle.masked_select(
  257. assigned_ltrb, bbox_mask).reshape([-1, 4])
  258. loss_dfl = self._df_loss(pred_dist_pos,
  259. assigned_ltrb_pos) * bbox_weight
  260. loss_dfl = loss_dfl.sum() / assigned_scores_sum
  261. else:
  262. loss_l1 = paddle.zeros([1])
  263. loss_iou = paddle.zeros([1])
  264. loss_dfl = pred_dist.sum() * 0.
  265. return loss_l1, loss_iou, loss_dfl
  266. def get_loss(self, head_outs, gt_meta):
  267. pred_scores, pred_distri, anchors,\
  268. anchor_points, num_anchors_list, stride_tensor = head_outs
  269. anchor_points_s = anchor_points / stride_tensor
  270. pred_bboxes = self._bbox_decode(anchor_points_s, pred_distri)
  271. gt_labels = gt_meta['gt_class']
  272. gt_bboxes = gt_meta['gt_bbox']
  273. pad_gt_mask = gt_meta['pad_gt_mask']
  274. # label assignment
  275. if gt_meta['epoch_id'] < self.static_assigner_epoch:
  276. assigned_labels, assigned_bboxes, assigned_scores = \
  277. self.static_assigner(
  278. anchors,
  279. num_anchors_list,
  280. gt_labels,
  281. gt_bboxes,
  282. pad_gt_mask,
  283. bg_index=self.num_classes,
  284. pred_bboxes=pred_bboxes.detach() * stride_tensor)
  285. alpha_l = 0.25
  286. else:
  287. assigned_labels, assigned_bboxes, assigned_scores = \
  288. self.assigner(
  289. pred_scores.detach(),
  290. pred_bboxes.detach() * stride_tensor,
  291. anchor_points,
  292. num_anchors_list,
  293. gt_labels,
  294. gt_bboxes,
  295. pad_gt_mask,
  296. bg_index=self.num_classes)
  297. alpha_l = -1
  298. # rescale bbox
  299. assigned_bboxes /= stride_tensor
  300. # cls loss
  301. if self.use_varifocal_loss:
  302. one_hot_label = F.one_hot(assigned_labels,
  303. self.num_classes + 1)[..., :-1]
  304. loss_cls = self._varifocal_loss(pred_scores, assigned_scores,
  305. one_hot_label)
  306. else:
  307. loss_cls = self._focal_loss(pred_scores, assigned_scores, alpha_l)
  308. assigned_scores_sum = assigned_scores.sum()
  309. if paddle.distributed.get_world_size() > 1:
  310. paddle.distributed.all_reduce(assigned_scores_sum)
  311. assigned_scores_sum = paddle.clip(
  312. assigned_scores_sum / paddle.distributed.get_world_size(),
  313. min=1)
  314. loss_cls /= assigned_scores_sum
  315. loss_l1, loss_iou, loss_dfl = \
  316. self._bbox_loss(pred_distri, pred_bboxes, anchor_points_s,
  317. assigned_labels, assigned_bboxes, assigned_scores,
  318. assigned_scores_sum)
  319. loss = self.loss_weight['class'] * loss_cls + \
  320. self.loss_weight['iou'] * loss_iou + \
  321. self.loss_weight['dfl'] * loss_dfl
  322. out_dict = {
  323. 'loss': loss,
  324. 'loss_cls': loss_cls,
  325. 'loss_iou': loss_iou,
  326. 'loss_dfl': loss_dfl,
  327. 'loss_l1': loss_l1,
  328. }
  329. return out_dict
  330. def post_process(self, head_outs, img_shape, scale_factor):
  331. pred_scores, pred_dist, anchor_points, stride_tensor = head_outs
  332. pred_bboxes = batch_distance2bbox(anchor_points,
  333. pred_dist.transpose([0, 2, 1]))
  334. pred_bboxes *= stride_tensor
  335. # scale bbox to origin
  336. scale_y, scale_x = paddle.split(scale_factor, 2, axis=-1)
  337. scale_factor = paddle.concat(
  338. [scale_x, scale_y, scale_x, scale_y], axis=-1).reshape([-1, 1, 4])
  339. pred_bboxes /= scale_factor
  340. if self.exclude_nms:
  341. # `exclude_nms=True` just use in benchmark
  342. return pred_bboxes.sum(), pred_scores.sum()
  343. else:
  344. bbox_pred, bbox_num, _ = self.nms(pred_bboxes, pred_scores)
  345. return bbox_pred, bbox_num