fcos_head.py 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270
  1. # Copyright (c) 2020 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. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import math
  18. import paddle
  19. import paddle.nn as nn
  20. import paddle.nn.functional as F
  21. from paddle import ParamAttr
  22. from paddle.nn.initializer import Normal, Constant
  23. from paddlex.ppdet.core.workspace import register
  24. from paddlex.ppdet.modeling.layers import ConvNormLayer
  25. class ScaleReg(nn.Layer):
  26. """
  27. Parameter for scaling the regression outputs.
  28. """
  29. def __init__(self):
  30. super(ScaleReg, self).__init__()
  31. self.scale_reg = self.create_parameter(
  32. shape=[1],
  33. attr=ParamAttr(initializer=Constant(value=1.)),
  34. dtype="float32")
  35. def forward(self, inputs):
  36. out = inputs * self.scale_reg
  37. return out
  38. @register
  39. class FCOSFeat(nn.Layer):
  40. """
  41. FCOSFeat of FCOS
  42. Args:
  43. feat_in (int): The channel number of input Tensor.
  44. feat_out (int): The channel number of output Tensor.
  45. num_convs (int): The convolution number of the FCOSFeat.
  46. norm_type (str): Normalization type, 'bn'/'sync_bn'/'gn'.
  47. use_dcn (bool): Whether to use dcn in tower or not.
  48. """
  49. def __init__(self,
  50. feat_in=256,
  51. feat_out=256,
  52. num_convs=4,
  53. norm_type='bn',
  54. use_dcn=False):
  55. super(FCOSFeat, self).__init__()
  56. self.num_convs = num_convs
  57. self.norm_type = norm_type
  58. self.cls_subnet_convs = []
  59. self.reg_subnet_convs = []
  60. for i in range(self.num_convs):
  61. in_c = feat_in if i == 0 else feat_out
  62. cls_conv_name = 'fcos_head_cls_tower_conv_{}'.format(i)
  63. cls_conv = self.add_sublayer(
  64. cls_conv_name,
  65. ConvNormLayer(
  66. ch_in=in_c,
  67. ch_out=feat_out,
  68. filter_size=3,
  69. stride=1,
  70. norm_type=norm_type,
  71. use_dcn=use_dcn,
  72. bias_on=True,
  73. lr_scale=2.))
  74. self.cls_subnet_convs.append(cls_conv)
  75. reg_conv_name = 'fcos_head_reg_tower_conv_{}'.format(i)
  76. reg_conv = self.add_sublayer(
  77. reg_conv_name,
  78. ConvNormLayer(
  79. ch_in=in_c,
  80. ch_out=feat_out,
  81. filter_size=3,
  82. stride=1,
  83. norm_type=norm_type,
  84. use_dcn=use_dcn,
  85. bias_on=True,
  86. lr_scale=2.))
  87. self.reg_subnet_convs.append(reg_conv)
  88. def forward(self, fpn_feat):
  89. cls_feat = fpn_feat
  90. reg_feat = fpn_feat
  91. for i in range(self.num_convs):
  92. cls_feat = F.relu(self.cls_subnet_convs[i](cls_feat))
  93. reg_feat = F.relu(self.reg_subnet_convs[i](reg_feat))
  94. return cls_feat, reg_feat
  95. @register
  96. class FCOSHead(nn.Layer):
  97. """
  98. FCOSHead
  99. Args:
  100. fcos_feat (object): Instance of 'FCOSFeat'
  101. num_classes (int): Number of classes
  102. fpn_stride (list): The stride of each FPN Layer
  103. prior_prob (float): Used to set the bias init for the class prediction layer
  104. fcos_loss (object): Instance of 'FCOSLoss'
  105. norm_reg_targets (bool): Normalization the regression target if true
  106. centerness_on_reg (bool): The prediction of centerness on regression or clssification branch
  107. """
  108. __inject__ = ['fcos_feat', 'fcos_loss']
  109. __shared__ = ['num_classes']
  110. def __init__(self,
  111. fcos_feat,
  112. num_classes=80,
  113. fpn_stride=[8, 16, 32, 64, 128],
  114. prior_prob=0.01,
  115. fcos_loss='FCOSLoss',
  116. norm_reg_targets=True,
  117. centerness_on_reg=True):
  118. super(FCOSHead, self).__init__()
  119. self.fcos_feat = fcos_feat
  120. self.num_classes = num_classes
  121. self.fpn_stride = fpn_stride
  122. self.prior_prob = prior_prob
  123. self.fcos_loss = fcos_loss
  124. self.norm_reg_targets = norm_reg_targets
  125. self.centerness_on_reg = centerness_on_reg
  126. conv_cls_name = "fcos_head_cls"
  127. bias_init_value = -math.log((1 - self.prior_prob) / self.prior_prob)
  128. self.fcos_head_cls = self.add_sublayer(
  129. conv_cls_name,
  130. nn.Conv2D(
  131. in_channels=256,
  132. out_channels=self.num_classes,
  133. kernel_size=3,
  134. stride=1,
  135. padding=1,
  136. weight_attr=ParamAttr(
  137. name=conv_cls_name + "_weights",
  138. initializer=Normal(
  139. mean=0., std=0.01)),
  140. bias_attr=ParamAttr(
  141. name=conv_cls_name + "_bias",
  142. initializer=Constant(value=bias_init_value))))
  143. conv_reg_name = "fcos_head_reg"
  144. self.fcos_head_reg = self.add_sublayer(
  145. conv_reg_name,
  146. nn.Conv2D(
  147. in_channels=256,
  148. out_channels=4,
  149. kernel_size=3,
  150. stride=1,
  151. padding=1,
  152. weight_attr=ParamAttr(
  153. name=conv_reg_name + "_weights",
  154. initializer=Normal(
  155. mean=0., std=0.01)),
  156. bias_attr=ParamAttr(
  157. name=conv_reg_name + "_bias",
  158. initializer=Constant(value=0))))
  159. conv_centerness_name = "fcos_head_centerness"
  160. self.fcos_head_centerness = self.add_sublayer(
  161. conv_centerness_name,
  162. nn.Conv2D(
  163. in_channels=256,
  164. out_channels=1,
  165. kernel_size=3,
  166. stride=1,
  167. padding=1,
  168. weight_attr=ParamAttr(
  169. name=conv_centerness_name + "_weights",
  170. initializer=Normal(
  171. mean=0., std=0.01)),
  172. bias_attr=ParamAttr(
  173. name=conv_centerness_name + "_bias",
  174. initializer=Constant(value=0))))
  175. self.scales_regs = []
  176. for i in range(len(self.fpn_stride)):
  177. lvl = int(math.log(int(self.fpn_stride[i]), 2))
  178. feat_name = 'p{}_feat'.format(lvl)
  179. scale_reg = self.add_sublayer(feat_name, ScaleReg())
  180. self.scales_regs.append(scale_reg)
  181. def _compute_locations_by_level(self, fpn_stride, feature):
  182. """
  183. Compute locations of anchor points of each FPN layer
  184. Args:
  185. fpn_stride (int): The stride of current FPN feature map
  186. feature (Tensor): Tensor of current FPN feature map
  187. Return:
  188. Anchor points locations of current FPN feature map
  189. """
  190. shape_fm = paddle.shape(feature)
  191. shape_fm.stop_gradient = True
  192. h, w = shape_fm[2], shape_fm[3]
  193. shift_x = paddle.arange(0, w * fpn_stride, fpn_stride)
  194. shift_y = paddle.arange(0, h * fpn_stride, fpn_stride)
  195. shift_x = paddle.unsqueeze(shift_x, axis=0)
  196. shift_y = paddle.unsqueeze(shift_y, axis=1)
  197. shift_x = paddle.expand(shift_x, shape=[h, w])
  198. shift_y = paddle.expand(shift_y, shape=[h, w])
  199. shift_x.stop_gradient = True
  200. shift_y.stop_gradient = True
  201. shift_x = paddle.reshape(shift_x, shape=[-1])
  202. shift_y = paddle.reshape(shift_y, shape=[-1])
  203. location = paddle.stack(
  204. [shift_x, shift_y], axis=-1) + float(fpn_stride) / 2
  205. location.stop_gradient = True
  206. return location
  207. def forward(self, fpn_feats, is_training):
  208. assert len(fpn_feats) == len(
  209. self.fpn_stride
  210. ), "The size of fpn_feats is not equal to size of fpn_stride"
  211. cls_logits_list = []
  212. bboxes_reg_list = []
  213. centerness_list = []
  214. for scale_reg, fpn_stride, fpn_feat in zip(self.scales_regs,
  215. self.fpn_stride, fpn_feats):
  216. fcos_cls_feat, fcos_reg_feat = self.fcos_feat(fpn_feat)
  217. cls_logits = self.fcos_head_cls(fcos_cls_feat)
  218. bbox_reg = scale_reg(self.fcos_head_reg(fcos_reg_feat))
  219. if self.centerness_on_reg:
  220. centerness = self.fcos_head_centerness(fcos_reg_feat)
  221. else:
  222. centerness = self.fcos_head_centerness(fcos_cls_feat)
  223. if self.norm_reg_targets:
  224. bbox_reg = F.relu(bbox_reg)
  225. if not is_training:
  226. bbox_reg = bbox_reg * fpn_stride
  227. else:
  228. bbox_reg = paddle.exp(bbox_reg)
  229. cls_logits_list.append(cls_logits)
  230. bboxes_reg_list.append(bbox_reg)
  231. centerness_list.append(centerness)
  232. if not is_training:
  233. locations_list = []
  234. for fpn_stride, feature in zip(self.fpn_stride, fpn_feats):
  235. location = self._compute_locations_by_level(fpn_stride,
  236. feature)
  237. locations_list.append(location)
  238. return locations_list, cls_logits_list, bboxes_reg_list, centerness_list
  239. else:
  240. return cls_logits_list, bboxes_reg_list, centerness_list
  241. def get_loss(self, fcos_head_outs, tag_labels, tag_bboxes, tag_centerness):
  242. cls_logits, bboxes_reg, centerness = fcos_head_outs
  243. return self.fcos_loss(cls_logits, bboxes_reg, centerness, tag_labels,
  244. tag_bboxes, tag_centerness)