darknet.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174
  1. # Copyright (c) 2019 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 six
  18. from paddle import fluid
  19. from paddle.fluid.param_attr import ParamAttr
  20. from paddle.fluid.regularizer import L2Decay
  21. from ppdet.core.workspace import register
  22. __all__ = ['DarkNet']
  23. @register
  24. class DarkNet(object):
  25. """
  26. DarkNet, see https://pjreddie.com/darknet/yolo/
  27. Args:
  28. depth (int): network depth, currently only darknet 53 is supported
  29. norm_type (str): normalization type, 'bn' and 'sync_bn' are supported
  30. norm_decay (float): weight decay for normalization layer weights
  31. """
  32. __shared__ = ['norm_type', 'weight_prefix_name']
  33. def __init__(self,
  34. depth=53,
  35. norm_type='bn',
  36. norm_decay=0.,
  37. weight_prefix_name='',
  38. freeze_at=-1):
  39. assert depth in [53], "unsupported depth value"
  40. self.depth = depth
  41. self.norm_type = norm_type
  42. self.norm_decay = norm_decay
  43. self.depth_cfg = {53: ([1, 2, 8, 8, 4], self.basicblock)}
  44. self.prefix_name = weight_prefix_name
  45. self.freeze_at = freeze_at
  46. def _conv_norm(self,
  47. input,
  48. ch_out,
  49. filter_size,
  50. stride,
  51. padding,
  52. act='leaky',
  53. name=None):
  54. conv = fluid.layers.conv2d(
  55. input=input,
  56. num_filters=ch_out,
  57. filter_size=filter_size,
  58. stride=stride,
  59. padding=padding,
  60. act=None,
  61. param_attr=ParamAttr(name=name + ".conv.weights"),
  62. bias_attr=False)
  63. bn_name = name + ".bn"
  64. bn_param_attr = ParamAttr(
  65. regularizer=L2Decay(float(self.norm_decay)),
  66. name=bn_name + '.scale')
  67. bn_bias_attr = ParamAttr(
  68. regularizer=L2Decay(float(self.norm_decay)),
  69. name=bn_name + '.offset')
  70. out = fluid.layers.batch_norm(
  71. input=conv,
  72. act=None,
  73. param_attr=bn_param_attr,
  74. bias_attr=bn_bias_attr,
  75. moving_mean_name=bn_name + '.mean',
  76. moving_variance_name=bn_name + '.var')
  77. # leaky relu here has `alpha` as 0.1, can not be set by
  78. # `act` param in fluid.layers.batch_norm above.
  79. if act == 'leaky':
  80. out = fluid.layers.leaky_relu(x=out, alpha=0.1)
  81. return out
  82. def _downsample(self,
  83. input,
  84. ch_out,
  85. filter_size=3,
  86. stride=2,
  87. padding=1,
  88. name=None):
  89. return self._conv_norm(
  90. input,
  91. ch_out=ch_out,
  92. filter_size=filter_size,
  93. stride=stride,
  94. padding=padding,
  95. name=name)
  96. def basicblock(self, input, ch_out, name=None):
  97. conv1 = self._conv_norm(
  98. input,
  99. ch_out=ch_out,
  100. filter_size=1,
  101. stride=1,
  102. padding=0,
  103. name=name + ".0")
  104. conv2 = self._conv_norm(
  105. conv1,
  106. ch_out=ch_out * 2,
  107. filter_size=3,
  108. stride=1,
  109. padding=1,
  110. name=name + ".1")
  111. out = fluid.layers.elementwise_add(x=input, y=conv2, act=None)
  112. return out
  113. def layer_warp(self, block_func, input, ch_out, count, name=None):
  114. out = block_func(input, ch_out=ch_out, name='{}.0'.format(name))
  115. for j in six.moves.xrange(1, count):
  116. out = block_func(out, ch_out=ch_out, name='{}.{}'.format(name, j))
  117. return out
  118. def __call__(self, input):
  119. """
  120. Get the backbone of DarkNet, that is output for the 5 stages.
  121. Args:
  122. input (Variable): input variable.
  123. Returns:
  124. The last variables of each stage.
  125. """
  126. stages, block_func = self.depth_cfg[self.depth]
  127. stages = stages[0:5]
  128. conv = self._conv_norm(
  129. input=input,
  130. ch_out=32,
  131. filter_size=3,
  132. stride=1,
  133. padding=1,
  134. name=self.prefix_name + "yolo_input")
  135. downsample_ = self._downsample(
  136. input=conv,
  137. ch_out=conv.shape[1] * 2,
  138. name=self.prefix_name + "yolo_input.downsample")
  139. blocks = []
  140. for i, stage in enumerate(stages):
  141. block = self.layer_warp(
  142. block_func=block_func,
  143. input=downsample_,
  144. ch_out=32 * 2**i,
  145. count=stage,
  146. name=self.prefix_name + "stage.{}".format(i))
  147. if i < self.freeze_at:
  148. block.stop_gradient = True
  149. blocks.append(block)
  150. if i < len(stages) - 1: # do not downsaple in the last stage
  151. downsample_ = self._downsample(
  152. input=block,
  153. ch_out=block.shape[1] * 2,
  154. name=self.prefix_name + "stage.{}.downsample".format(i))
  155. return blocks