custom_pan.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227
  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 as nn
  16. import paddle.nn.functional as F
  17. from paddlex.ppdet.core.workspace import register, serializable
  18. from paddlex.ppdet.modeling.layers import DropBlock
  19. from paddlex.ppdet.modeling.ops import get_act_fn
  20. from ..backbones.cspresnet import ConvBNLayer, BasicBlock
  21. from ..shape_spec import ShapeSpec
  22. __all__ = ['CustomCSPPAN']
  23. class SPP(nn.Layer):
  24. def __init__(self,
  25. ch_in,
  26. ch_out,
  27. k,
  28. pool_size,
  29. act='swish',
  30. data_format='NCHW'):
  31. super(SPP, self).__init__()
  32. self.pool = []
  33. self.data_format = data_format
  34. for i, size in enumerate(pool_size):
  35. pool = self.add_sublayer(
  36. 'pool{}'.format(i),
  37. nn.MaxPool2D(
  38. kernel_size=size,
  39. stride=1,
  40. padding=size // 2,
  41. data_format=data_format,
  42. ceil_mode=False))
  43. self.pool.append(pool)
  44. self.conv = ConvBNLayer(ch_in, ch_out, k, padding=k // 2, act=act)
  45. def forward(self, x):
  46. outs = [x]
  47. for pool in self.pool:
  48. outs.append(pool(x))
  49. if self.data_format == 'NCHW':
  50. y = paddle.concat(outs, axis=1)
  51. else:
  52. y = paddle.concat(outs, axis=-1)
  53. y = self.conv(y)
  54. return y
  55. class CSPStage(nn.Layer):
  56. def __init__(self, block_fn, ch_in, ch_out, n, act='swish', spp=False):
  57. super(CSPStage, self).__init__()
  58. ch_mid = int(ch_out // 2)
  59. self.conv1 = ConvBNLayer(ch_in, ch_mid, 1, act=act)
  60. self.conv2 = ConvBNLayer(ch_in, ch_mid, 1, act=act)
  61. self.convs = nn.Sequential()
  62. next_ch_in = ch_mid
  63. for i in range(n):
  64. self.convs.add_sublayer(
  65. str(i),
  66. eval(block_fn)(next_ch_in, ch_mid, act=act, shortcut=False))
  67. if i == (n - 1) // 2 and spp:
  68. self.convs.add_sublayer(
  69. 'spp', SPP(ch_mid * 4, ch_mid, 1, [5, 9, 13], act=act))
  70. next_ch_in = ch_mid
  71. self.conv3 = ConvBNLayer(ch_mid * 2, ch_out, 1, act=act)
  72. def forward(self, x):
  73. y1 = self.conv1(x)
  74. y2 = self.conv2(x)
  75. y2 = self.convs(y2)
  76. y = paddle.concat([y1, y2], axis=1)
  77. y = self.conv3(y)
  78. return y
  79. @register
  80. @serializable
  81. class CustomCSPPAN(nn.Layer):
  82. __shared__ = [
  83. 'norm_type', 'data_format', 'width_mult', 'depth_mult', 'trt'
  84. ]
  85. def __init__(self,
  86. in_channels=[256, 512, 1024],
  87. out_channels=[1024, 512, 256],
  88. norm_type='bn',
  89. act='leaky',
  90. stage_fn='CSPStage',
  91. block_fn='BasicBlock',
  92. stage_num=1,
  93. block_num=3,
  94. drop_block=False,
  95. block_size=3,
  96. keep_prob=0.9,
  97. spp=False,
  98. data_format='NCHW',
  99. width_mult=1.0,
  100. depth_mult=1.0,
  101. trt=False):
  102. super(CustomCSPPAN, self).__init__()
  103. out_channels = [max(round(c * width_mult), 1) for c in out_channels]
  104. block_num = max(round(block_num * depth_mult), 1)
  105. act = get_act_fn(
  106. act, trt=trt) if act is None or isinstance(act,
  107. (str, dict)) else act
  108. self.num_blocks = len(in_channels)
  109. self.data_format = data_format
  110. self._out_channels = out_channels
  111. in_channels = in_channels[::-1]
  112. fpn_stages = []
  113. fpn_routes = []
  114. for i, (ch_in, ch_out) in enumerate(zip(in_channels, out_channels)):
  115. if i > 0:
  116. ch_in += ch_pre // 2
  117. stage = nn.Sequential()
  118. for j in range(stage_num):
  119. stage.add_sublayer(
  120. str(j),
  121. eval(stage_fn)(block_fn,
  122. ch_in if j == 0 else ch_out,
  123. ch_out,
  124. block_num,
  125. act=act,
  126. spp=(spp and i == 0)))
  127. if drop_block:
  128. stage.add_sublayer('drop', DropBlock(block_size, keep_prob))
  129. fpn_stages.append(stage)
  130. if i < self.num_blocks - 1:
  131. fpn_routes.append(
  132. ConvBNLayer(
  133. ch_in=ch_out,
  134. ch_out=ch_out // 2,
  135. filter_size=1,
  136. stride=1,
  137. padding=0,
  138. act=act))
  139. ch_pre = ch_out
  140. self.fpn_stages = nn.LayerList(fpn_stages)
  141. self.fpn_routes = nn.LayerList(fpn_routes)
  142. pan_stages = []
  143. pan_routes = []
  144. for i in reversed(range(self.num_blocks - 1)):
  145. pan_routes.append(
  146. ConvBNLayer(
  147. ch_in=out_channels[i + 1],
  148. ch_out=out_channels[i + 1],
  149. filter_size=3,
  150. stride=2,
  151. padding=1,
  152. act=act))
  153. ch_in = out_channels[i] + out_channels[i + 1]
  154. ch_out = out_channels[i]
  155. stage = nn.Sequential()
  156. for j in range(stage_num):
  157. stage.add_sublayer(
  158. str(j),
  159. eval(stage_fn)(block_fn,
  160. ch_in if j == 0 else ch_out,
  161. ch_out,
  162. block_num,
  163. act=act,
  164. spp=False))
  165. if drop_block:
  166. stage.add_sublayer('drop', DropBlock(block_size, keep_prob))
  167. pan_stages.append(stage)
  168. self.pan_stages = nn.LayerList(pan_stages[::-1])
  169. self.pan_routes = nn.LayerList(pan_routes[::-1])
  170. def forward(self, blocks, for_mot=False):
  171. blocks = blocks[::-1]
  172. fpn_feats = []
  173. for i, block in enumerate(blocks):
  174. if i > 0:
  175. block = paddle.concat([route, block], axis=1)
  176. route = self.fpn_stages[i](block)
  177. fpn_feats.append(route)
  178. if i < self.num_blocks - 1:
  179. route = self.fpn_routes[i](route)
  180. route = F.interpolate(
  181. route, scale_factor=2., data_format=self.data_format)
  182. pan_feats = [fpn_feats[-1], ]
  183. route = fpn_feats[-1]
  184. for i in reversed(range(self.num_blocks - 1)):
  185. block = fpn_feats[i]
  186. route = self.pan_routes[i](route)
  187. block = paddle.concat([route, block], axis=1)
  188. route = self.pan_stages[i](block)
  189. pan_feats.append(route)
  190. return pan_feats[::-1]
  191. @classmethod
  192. def from_config(cls, cfg, input_shape):
  193. return {'in_channels': [i.channels for i in input_shape], }
  194. @property
  195. def out_shape(self):
  196. return [ShapeSpec(channels=c) for c in self._out_channels]