ssd_head.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  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. class SepConvLayer(nn.Layer):
  22. def __init__(self,
  23. in_channels,
  24. out_channels,
  25. kernel_size=3,
  26. padding=1,
  27. conv_decay=0):
  28. super(SepConvLayer, self).__init__()
  29. self.dw_conv = nn.Conv2D(
  30. in_channels=in_channels,
  31. out_channels=in_channels,
  32. kernel_size=kernel_size,
  33. stride=1,
  34. padding=padding,
  35. groups=in_channels,
  36. weight_attr=ParamAttr(regularizer=L2Decay(conv_decay)),
  37. bias_attr=False)
  38. self.bn = nn.BatchNorm2D(
  39. in_channels,
  40. weight_attr=ParamAttr(regularizer=L2Decay(0.)),
  41. bias_attr=ParamAttr(regularizer=L2Decay(0.)))
  42. self.pw_conv = nn.Conv2D(
  43. in_channels=in_channels,
  44. out_channels=out_channels,
  45. kernel_size=1,
  46. stride=1,
  47. padding=0,
  48. weight_attr=ParamAttr(regularizer=L2Decay(conv_decay)),
  49. bias_attr=False)
  50. def forward(self, x):
  51. x = self.dw_conv(x)
  52. x = F.relu6(self.bn(x))
  53. x = self.pw_conv(x)
  54. return x
  55. @register
  56. class SSDHead(nn.Layer):
  57. """
  58. SSDHead
  59. Args:
  60. num_classes (int): Number of classes
  61. in_channels (list): Number of channels per input feature
  62. anchor_generator (dict): Configuration of 'AnchorGeneratorSSD' instance
  63. kernel_size (int): Conv kernel size
  64. padding (int): Conv padding
  65. use_sepconv (bool): Use SepConvLayer if true
  66. conv_decay (float): Conv regularization coeff
  67. loss (object): 'SSDLoss' instance
  68. """
  69. __shared__ = ['num_classes']
  70. __inject__ = ['anchor_generator', 'loss']
  71. def __init__(self,
  72. num_classes=80,
  73. in_channels=(512, 1024, 512, 256, 256, 256),
  74. anchor_generator=AnchorGeneratorSSD().__dict__,
  75. kernel_size=3,
  76. padding=1,
  77. use_sepconv=False,
  78. conv_decay=0.,
  79. loss='SSDLoss'):
  80. super(SSDHead, self).__init__()
  81. # add background class
  82. self.num_classes = num_classes + 1
  83. self.in_channels = in_channels
  84. self.anchor_generator = anchor_generator
  85. self.loss = loss
  86. if isinstance(anchor_generator, dict):
  87. self.anchor_generator = AnchorGeneratorSSD(**anchor_generator)
  88. self.num_priors = self.anchor_generator.num_priors
  89. self.box_convs = []
  90. self.score_convs = []
  91. for i, num_prior in enumerate(self.num_priors):
  92. box_conv_name = "boxes{}".format(i)
  93. if not use_sepconv:
  94. box_conv = self.add_sublayer(
  95. box_conv_name,
  96. nn.Conv2D(
  97. in_channels=in_channels[i],
  98. out_channels=num_prior * 4,
  99. kernel_size=kernel_size,
  100. padding=padding))
  101. else:
  102. box_conv = self.add_sublayer(
  103. box_conv_name,
  104. SepConvLayer(
  105. in_channels=in_channels[i],
  106. out_channels=num_prior * 4,
  107. kernel_size=kernel_size,
  108. padding=padding,
  109. conv_decay=conv_decay))
  110. self.box_convs.append(box_conv)
  111. score_conv_name = "scores{}".format(i)
  112. if not use_sepconv:
  113. score_conv = self.add_sublayer(
  114. score_conv_name,
  115. nn.Conv2D(
  116. in_channels=in_channels[i],
  117. out_channels=num_prior * self.num_classes,
  118. kernel_size=kernel_size,
  119. padding=padding))
  120. else:
  121. score_conv = self.add_sublayer(
  122. score_conv_name,
  123. SepConvLayer(
  124. in_channels=in_channels[i],
  125. out_channels=num_prior * self.num_classes,
  126. kernel_size=kernel_size,
  127. padding=padding,
  128. conv_decay=conv_decay))
  129. self.score_convs.append(score_conv)
  130. @classmethod
  131. def from_config(cls, cfg, input_shape):
  132. return {'in_channels': [i.channels for i in input_shape], }
  133. def forward(self, feats, image, gt_bbox=None, gt_class=None):
  134. box_preds = []
  135. cls_scores = []
  136. prior_boxes = []
  137. for feat, box_conv, score_conv in zip(feats, self.box_convs,
  138. self.score_convs):
  139. box_pred = box_conv(feat)
  140. box_pred = paddle.transpose(box_pred, [0, 2, 3, 1])
  141. box_pred = paddle.reshape(box_pred, [0, -1, 4])
  142. box_preds.append(box_pred)
  143. cls_score = score_conv(feat)
  144. cls_score = paddle.transpose(cls_score, [0, 2, 3, 1])
  145. cls_score = paddle.reshape(cls_score, [0, -1, self.num_classes])
  146. cls_scores.append(cls_score)
  147. prior_boxes = self.anchor_generator(feats, image)
  148. if self.training:
  149. return self.get_loss(box_preds, cls_scores, gt_bbox, gt_class,
  150. prior_boxes)
  151. else:
  152. return (box_preds, cls_scores), prior_boxes
  153. def get_loss(self, boxes, scores, gt_bbox, gt_class, prior_boxes):
  154. return self.loss(boxes, scores, gt_bbox, gt_class, prior_boxes)