blazeface_fpn.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216
  1. # Copyright (c) 2021 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.functional as F
  16. from paddle import ParamAttr
  17. import paddle.nn as nn
  18. from paddle.nn.initializer import KaimingNormal
  19. from paddlex.ppdet.core.workspace import register, serializable
  20. from ..shape_spec import ShapeSpec
  21. __all__ = ['BlazeNeck']
  22. def hard_swish(x):
  23. return x * F.relu6(x + 3) / 6.
  24. class ConvBNLayer(nn.Layer):
  25. def __init__(self,
  26. in_channels,
  27. out_channels,
  28. kernel_size,
  29. stride,
  30. padding,
  31. num_groups=1,
  32. act='relu',
  33. conv_lr=0.1,
  34. conv_decay=0.,
  35. norm_decay=0.,
  36. norm_type='bn',
  37. name=None):
  38. super(ConvBNLayer, self).__init__()
  39. self.act = act
  40. self._conv = nn.Conv2D(
  41. in_channels,
  42. out_channels,
  43. kernel_size=kernel_size,
  44. stride=stride,
  45. padding=padding,
  46. groups=num_groups,
  47. weight_attr=ParamAttr(
  48. learning_rate=conv_lr, initializer=KaimingNormal()),
  49. bias_attr=False)
  50. if norm_type == 'sync_bn':
  51. self._batch_norm = nn.SyncBatchNorm(out_channels)
  52. else:
  53. self._batch_norm = nn.BatchNorm(
  54. out_channels, act=None, use_global_stats=False)
  55. def forward(self, x):
  56. x = self._conv(x)
  57. x = self._batch_norm(x)
  58. if self.act == "relu":
  59. x = F.relu(x)
  60. elif self.act == "relu6":
  61. x = F.relu6(x)
  62. elif self.act == 'leaky':
  63. x = F.leaky_relu(x)
  64. elif self.act == 'hard_swish':
  65. x = hard_swish(x)
  66. return x
  67. class FPN(nn.Layer):
  68. def __init__(self, in_channels, out_channels, name=None):
  69. super(FPN, self).__init__()
  70. self.conv1_fpn = ConvBNLayer(
  71. in_channels,
  72. out_channels // 2,
  73. kernel_size=1,
  74. padding=0,
  75. stride=1,
  76. act='leaky',
  77. name=name + '_output1')
  78. self.conv2_fpn = ConvBNLayer(
  79. in_channels,
  80. out_channels // 2,
  81. kernel_size=1,
  82. padding=0,
  83. stride=1,
  84. act='leaky',
  85. name=name + '_output2')
  86. self.conv3_fpn = ConvBNLayer(
  87. out_channels // 2,
  88. out_channels // 2,
  89. kernel_size=3,
  90. padding=1,
  91. stride=1,
  92. act='leaky',
  93. name=name + '_merge')
  94. def forward(self, input):
  95. output1 = self.conv1_fpn(input[0])
  96. output2 = self.conv2_fpn(input[1])
  97. up2 = F.upsample(
  98. output2, size=paddle.shape(output1)[-2:], mode='nearest')
  99. output1 = paddle.add(output1, up2)
  100. output1 = self.conv3_fpn(output1)
  101. return output1, output2
  102. class SSH(nn.Layer):
  103. def __init__(self, in_channels, out_channels, name=None):
  104. super(SSH, self).__init__()
  105. assert out_channels % 4 == 0
  106. self.conv0_ssh = ConvBNLayer(
  107. in_channels,
  108. out_channels // 2,
  109. kernel_size=3,
  110. padding=1,
  111. stride=1,
  112. act=None,
  113. name=name + 'ssh_conv3')
  114. self.conv1_ssh = ConvBNLayer(
  115. out_channels // 2,
  116. out_channels // 4,
  117. kernel_size=3,
  118. padding=1,
  119. stride=1,
  120. act='leaky',
  121. name=name + 'ssh_conv5_1')
  122. self.conv2_ssh = ConvBNLayer(
  123. out_channels // 4,
  124. out_channels // 4,
  125. kernel_size=3,
  126. padding=1,
  127. stride=1,
  128. act=None,
  129. name=name + 'ssh_conv5_2')
  130. self.conv3_ssh = ConvBNLayer(
  131. out_channels // 4,
  132. out_channels // 4,
  133. kernel_size=3,
  134. padding=1,
  135. stride=1,
  136. act='leaky',
  137. name=name + 'ssh_conv7_1')
  138. self.conv4_ssh = ConvBNLayer(
  139. out_channels // 4,
  140. out_channels // 4,
  141. kernel_size=3,
  142. padding=1,
  143. stride=1,
  144. act=None,
  145. name=name + 'ssh_conv7_2')
  146. def forward(self, x):
  147. conv0 = self.conv0_ssh(x)
  148. conv1 = self.conv1_ssh(conv0)
  149. conv2 = self.conv2_ssh(conv1)
  150. conv3 = self.conv3_ssh(conv2)
  151. conv4 = self.conv4_ssh(conv3)
  152. concat = paddle.concat([conv0, conv2, conv4], axis=1)
  153. return F.relu(concat)
  154. @register
  155. @serializable
  156. class BlazeNeck(nn.Layer):
  157. def __init__(self, in_channel, neck_type="None", data_format='NCHW'):
  158. super(BlazeNeck, self).__init__()
  159. self.neck_type = neck_type
  160. self.reture_input = False
  161. self._out_channels = in_channel
  162. if self.neck_type == 'None':
  163. self.reture_input = True
  164. if "fpn" in self.neck_type:
  165. self.fpn = FPN(self._out_channels[0],
  166. self._out_channels[1],
  167. name='fpn')
  168. self._out_channels = [
  169. self._out_channels[0] // 2, self._out_channels[1] // 2
  170. ]
  171. if "ssh" in self.neck_type:
  172. self.ssh1 = SSH(self._out_channels[0],
  173. self._out_channels[0],
  174. name='ssh1')
  175. self.ssh2 = SSH(self._out_channels[1],
  176. self._out_channels[1],
  177. name='ssh2')
  178. self._out_channels = [self._out_channels[0], self._out_channels[1]]
  179. def forward(self, inputs):
  180. if self.reture_input:
  181. return inputs
  182. output1, output2 = None, None
  183. if "fpn" in self.neck_type:
  184. backout_4, backout_1 = inputs
  185. output1, output2 = self.fpn([backout_4, backout_1])
  186. if self.neck_type == "only_fpn":
  187. return [output1, output2]
  188. if self.neck_type == "only_ssh":
  189. output1, output2 = inputs
  190. feature1 = self.ssh1(output1)
  191. feature2 = self.ssh2(output2)
  192. return [feature1, feature2]
  193. @property
  194. def out_shape(self):
  195. return [
  196. ShapeSpec(channels=c)
  197. for c in [self._out_channels[0], self._out_channels[1]]
  198. ]