ssd_head.py 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217
  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. 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 paddle.regularizer import L2Decay
  19. from paddle import ParamAttr
  20. from ..layers import AnchorGeneratorSSD
  21. from ..cls_utils import _get_class_default_kwargs
  22. class SepConvLayer(nn.Layer):
  23. def __init__(self,
  24. in_channels,
  25. out_channels,
  26. kernel_size=3,
  27. padding=1,
  28. conv_decay=0.):
  29. super(SepConvLayer, self).__init__()
  30. self.dw_conv = nn.Conv2D(
  31. in_channels=in_channels,
  32. out_channels=in_channels,
  33. kernel_size=kernel_size,
  34. stride=1,
  35. padding=padding,
  36. groups=in_channels,
  37. weight_attr=ParamAttr(regularizer=L2Decay(conv_decay)),
  38. bias_attr=False)
  39. self.bn = nn.BatchNorm2D(
  40. in_channels,
  41. weight_attr=ParamAttr(regularizer=L2Decay(0.)),
  42. bias_attr=ParamAttr(regularizer=L2Decay(0.)))
  43. self.pw_conv = nn.Conv2D(
  44. in_channels=in_channels,
  45. out_channels=out_channels,
  46. kernel_size=1,
  47. stride=1,
  48. padding=0,
  49. weight_attr=ParamAttr(regularizer=L2Decay(conv_decay)),
  50. bias_attr=False)
  51. def forward(self, x):
  52. x = self.dw_conv(x)
  53. x = F.relu6(self.bn(x))
  54. x = self.pw_conv(x)
  55. return x
  56. class SSDExtraHead(nn.Layer):
  57. def __init__(self,
  58. in_channels=256,
  59. out_channels=([256, 512], [256, 512], [128, 256], [128, 256],
  60. [128, 256]),
  61. strides=(2, 2, 2, 1, 1),
  62. paddings=(1, 1, 1, 0, 0)):
  63. super(SSDExtraHead, self).__init__()
  64. self.convs = nn.LayerList()
  65. for out_channel, stride, padding in zip(out_channels, strides,
  66. paddings):
  67. self.convs.append(
  68. self._make_layers(in_channels, out_channel[0], out_channel[1],
  69. stride, padding))
  70. in_channels = out_channel[-1]
  71. def _make_layers(self, c_in, c_hidden, c_out, stride_3x3, padding_3x3):
  72. return nn.Sequential(
  73. nn.Conv2D(c_in, c_hidden, 1),
  74. nn.ReLU(),
  75. nn.Conv2D(c_hidden, c_out, 3, stride_3x3, padding_3x3), nn.ReLU())
  76. def forward(self, x):
  77. out = [x]
  78. for conv_layer in self.convs:
  79. out.append(conv_layer(out[-1]))
  80. return out
  81. @register
  82. class SSDHead(nn.Layer):
  83. """
  84. SSDHead
  85. Args:
  86. num_classes (int): Number of classes
  87. in_channels (list): Number of channels per input feature
  88. anchor_generator (dict): Configuration of 'AnchorGeneratorSSD' instance
  89. kernel_size (int): Conv kernel size
  90. padding (int): Conv padding
  91. use_sepconv (bool): Use SepConvLayer if true
  92. conv_decay (float): Conv regularization coeff
  93. loss (object): 'SSDLoss' instance
  94. use_extra_head (bool): If use ResNet34 as baskbone, you should set `use_extra_head`=True
  95. """
  96. __shared__ = ['num_classes']
  97. __inject__ = ['anchor_generator', 'loss']
  98. def __init__(
  99. self,
  100. num_classes=80,
  101. in_channels=(512, 1024, 512, 256, 256, 256),
  102. anchor_generator=_get_class_default_kwargs(AnchorGeneratorSSD),
  103. kernel_size=3,
  104. padding=1,
  105. use_sepconv=False,
  106. conv_decay=0.,
  107. loss='SSDLoss',
  108. use_extra_head=False):
  109. super(SSDHead, self).__init__()
  110. # add background class
  111. self.num_classes = num_classes + 1
  112. self.in_channels = in_channels
  113. self.anchor_generator = anchor_generator
  114. self.loss = loss
  115. self.use_extra_head = use_extra_head
  116. if self.use_extra_head:
  117. self.ssd_extra_head = SSDExtraHead()
  118. self.in_channels = [256, 512, 512, 256, 256, 256]
  119. if isinstance(anchor_generator, dict):
  120. self.anchor_generator = AnchorGeneratorSSD(**anchor_generator)
  121. self.num_priors = self.anchor_generator.num_priors
  122. self.box_convs = []
  123. self.score_convs = []
  124. for i, num_prior in enumerate(self.num_priors):
  125. box_conv_name = "boxes{}".format(i)
  126. if not use_sepconv:
  127. box_conv = self.add_sublayer(
  128. box_conv_name,
  129. nn.Conv2D(
  130. in_channels=self.in_channels[i],
  131. out_channels=num_prior * 4,
  132. kernel_size=kernel_size,
  133. padding=padding))
  134. else:
  135. box_conv = self.add_sublayer(
  136. box_conv_name,
  137. SepConvLayer(
  138. in_channels=self.in_channels[i],
  139. out_channels=num_prior * 4,
  140. kernel_size=kernel_size,
  141. padding=padding,
  142. conv_decay=conv_decay))
  143. self.box_convs.append(box_conv)
  144. score_conv_name = "scores{}".format(i)
  145. if not use_sepconv:
  146. score_conv = self.add_sublayer(
  147. score_conv_name,
  148. nn.Conv2D(
  149. in_channels=self.in_channels[i],
  150. out_channels=num_prior * self.num_classes,
  151. kernel_size=kernel_size,
  152. padding=padding))
  153. else:
  154. score_conv = self.add_sublayer(
  155. score_conv_name,
  156. SepConvLayer(
  157. in_channels=self.in_channels[i],
  158. out_channels=num_prior * self.num_classes,
  159. kernel_size=kernel_size,
  160. padding=padding,
  161. conv_decay=conv_decay))
  162. self.score_convs.append(score_conv)
  163. @classmethod
  164. def from_config(cls, cfg, input_shape):
  165. return {'in_channels': [i.channels for i in input_shape], }
  166. def forward(self, feats, image, gt_bbox=None, gt_class=None):
  167. if self.use_extra_head:
  168. assert len(feats) == 1, \
  169. ("If you set use_extra_head=True, backbone feature "
  170. "list length should be 1.")
  171. feats = self.ssd_extra_head(feats[0])
  172. box_preds = []
  173. cls_scores = []
  174. for feat, box_conv, score_conv in zip(feats, self.box_convs,
  175. self.score_convs):
  176. box_pred = box_conv(feat)
  177. box_pred = paddle.transpose(box_pred, [0, 2, 3, 1])
  178. box_pred = paddle.reshape(box_pred, [0, -1, 4])
  179. box_preds.append(box_pred)
  180. cls_score = score_conv(feat)
  181. cls_score = paddle.transpose(cls_score, [0, 2, 3, 1])
  182. cls_score = paddle.reshape(cls_score, [0, -1, self.num_classes])
  183. cls_scores.append(cls_score)
  184. prior_boxes = self.anchor_generator(feats, image)
  185. if self.training:
  186. return self.get_loss(box_preds, cls_scores, gt_bbox, gt_class,
  187. prior_boxes)
  188. else:
  189. return (box_preds, cls_scores), prior_boxes
  190. def get_loss(self, boxes, scores, gt_bbox, gt_class, prior_boxes):
  191. return self.loss(boxes, scores, gt_bbox, gt_class, prior_boxes)