cspresnet.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283
  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. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import paddle
  18. import paddle.nn as nn
  19. import paddle.nn.functional as F
  20. from paddle import ParamAttr
  21. from paddle.regularizer import L2Decay
  22. from paddlex.ppdet.modeling.ops import get_act_fn
  23. from paddlex.ppdet.core.workspace import register, serializable
  24. from ..shape_spec import ShapeSpec
  25. __all__ = ['CSPResNet', 'BasicBlock', 'EffectiveSELayer', 'ConvBNLayer']
  26. class ConvBNLayer(nn.Layer):
  27. def __init__(self,
  28. ch_in,
  29. ch_out,
  30. filter_size=3,
  31. stride=1,
  32. groups=1,
  33. padding=0,
  34. act=None):
  35. super(ConvBNLayer, self).__init__()
  36. self.conv = nn.Conv2D(
  37. in_channels=ch_in,
  38. out_channels=ch_out,
  39. kernel_size=filter_size,
  40. stride=stride,
  41. padding=padding,
  42. groups=groups,
  43. bias_attr=False)
  44. self.bn = nn.BatchNorm2D(
  45. ch_out,
  46. weight_attr=ParamAttr(regularizer=L2Decay(0.0)),
  47. bias_attr=ParamAttr(regularizer=L2Decay(0.0)))
  48. self.act = get_act_fn(act) if act is None or isinstance(act, (
  49. str, dict)) else act
  50. def forward(self, x):
  51. x = self.conv(x)
  52. x = self.bn(x)
  53. x = self.act(x)
  54. return x
  55. class RepVggBlock(nn.Layer):
  56. def __init__(self, ch_in, ch_out, act='relu'):
  57. super(RepVggBlock, self).__init__()
  58. self.ch_in = ch_in
  59. self.ch_out = ch_out
  60. self.conv1 = ConvBNLayer(
  61. ch_in, ch_out, 3, stride=1, padding=1, act=None)
  62. self.conv2 = ConvBNLayer(
  63. ch_in, ch_out, 1, stride=1, padding=0, act=None)
  64. self.act = get_act_fn(act) if act is None or isinstance(act, (
  65. str, dict)) else act
  66. def forward(self, x):
  67. if hasattr(self, 'conv'):
  68. y = self.conv(x)
  69. else:
  70. y = self.conv1(x) + self.conv2(x)
  71. y = self.act(y)
  72. return y
  73. def convert_to_deploy(self):
  74. if not hasattr(self, 'conv'):
  75. self.conv = nn.Conv2D(
  76. in_channels=self.ch_in,
  77. out_channels=self.ch_out,
  78. kernel_size=3,
  79. stride=1,
  80. padding=1,
  81. groups=1)
  82. kernel, bias = self.get_equivalent_kernel_bias()
  83. self.conv.weight.set_value(kernel)
  84. self.conv.bias.set_value(bias)
  85. self.__delattr__('conv1')
  86. self.__delattr__('conv2')
  87. def get_equivalent_kernel_bias(self):
  88. kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)
  89. kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)
  90. return kernel3x3 + self._pad_1x1_to_3x3_tensor(
  91. kernel1x1), bias3x3 + bias1x1
  92. def _pad_1x1_to_3x3_tensor(self, kernel1x1):
  93. if kernel1x1 is None:
  94. return 0
  95. else:
  96. return nn.functional.pad(kernel1x1, [1, 1, 1, 1])
  97. def _fuse_bn_tensor(self, branch):
  98. if branch is None:
  99. return 0, 0
  100. kernel = branch.conv.weight
  101. running_mean = branch.bn._mean
  102. running_var = branch.bn._variance
  103. gamma = branch.bn.weight
  104. beta = branch.bn.bias
  105. eps = branch.bn._epsilon
  106. std = (running_var + eps).sqrt()
  107. t = (gamma / std).reshape((-1, 1, 1, 1))
  108. return kernel * t, beta - running_mean * gamma / std
  109. class BasicBlock(nn.Layer):
  110. def __init__(self, ch_in, ch_out, act='relu', shortcut=True):
  111. super(BasicBlock, self).__init__()
  112. assert ch_in == ch_out
  113. self.conv1 = ConvBNLayer(
  114. ch_in, ch_out, 3, stride=1, padding=1, act=act)
  115. self.conv2 = RepVggBlock(ch_out, ch_out, act=act)
  116. self.shortcut = shortcut
  117. def forward(self, x):
  118. y = self.conv1(x)
  119. y = self.conv2(y)
  120. if self.shortcut:
  121. return paddle.add(x, y)
  122. else:
  123. return y
  124. class EffectiveSELayer(nn.Layer):
  125. """ Effective Squeeze-Excitation
  126. From `CenterMask : Real-Time Anchor-Free Instance Segmentation` - https://arxiv.org/abs/1911.06667
  127. """
  128. def __init__(self, channels, act='hardsigmoid'):
  129. super(EffectiveSELayer, self).__init__()
  130. self.fc = nn.Conv2D(channels, channels, kernel_size=1, padding=0)
  131. self.act = get_act_fn(act) if act is None or isinstance(act, (
  132. str, dict)) else act
  133. def forward(self, x):
  134. x_se = x.mean((2, 3), keepdim=True)
  135. x_se = self.fc(x_se)
  136. return x * self.act(x_se)
  137. class CSPResStage(nn.Layer):
  138. def __init__(self,
  139. block_fn,
  140. ch_in,
  141. ch_out,
  142. n,
  143. stride,
  144. act='relu',
  145. attn='eca'):
  146. super(CSPResStage, self).__init__()
  147. ch_mid = (ch_in + ch_out) // 2
  148. if stride == 2:
  149. self.conv_down = ConvBNLayer(
  150. ch_in, ch_mid, 3, stride=2, padding=1, act=act)
  151. else:
  152. self.conv_down = None
  153. self.conv1 = ConvBNLayer(ch_mid, ch_mid // 2, 1, act=act)
  154. self.conv2 = ConvBNLayer(ch_mid, ch_mid // 2, 1, act=act)
  155. self.blocks = nn.Sequential(* [
  156. block_fn(
  157. ch_mid // 2, ch_mid // 2, act=act, shortcut=True)
  158. for i in range(n)
  159. ])
  160. if attn:
  161. self.attn = EffectiveSELayer(ch_mid, act='hardsigmoid')
  162. else:
  163. self.attn = None
  164. self.conv3 = ConvBNLayer(ch_mid, ch_out, 1, act=act)
  165. def forward(self, x):
  166. if self.conv_down is not None:
  167. x = self.conv_down(x)
  168. y1 = self.conv1(x)
  169. y2 = self.blocks(self.conv2(x))
  170. y = paddle.concat([y1, y2], axis=1)
  171. if self.attn is not None:
  172. y = self.attn(y)
  173. y = self.conv3(y)
  174. return y
  175. @register
  176. @serializable
  177. class CSPResNet(nn.Layer):
  178. __shared__ = ['width_mult', 'depth_mult', 'trt']
  179. def __init__(self,
  180. layers=[3, 6, 6, 3],
  181. channels=[64, 128, 256, 512, 1024],
  182. act='swish',
  183. return_idx=[0, 1, 2, 3, 4],
  184. depth_wise=False,
  185. use_large_stem=False,
  186. width_mult=1.0,
  187. depth_mult=1.0,
  188. trt=False):
  189. super(CSPResNet, self).__init__()
  190. channels = [max(round(c * width_mult), 1) for c in channels]
  191. layers = [max(round(l * depth_mult), 1) for l in layers]
  192. act = get_act_fn(
  193. act, trt=trt) if act is None or isinstance(act,
  194. (str, dict)) else act
  195. if use_large_stem:
  196. self.stem = nn.Sequential(
  197. ('conv1', ConvBNLayer(
  198. 3, channels[0] // 2, 3, stride=2, padding=1, act=act)),
  199. ('conv2', ConvBNLayer(
  200. channels[0] // 2,
  201. channels[0] // 2,
  202. 3,
  203. stride=1,
  204. padding=1,
  205. act=act)), ('conv3', ConvBNLayer(
  206. channels[0] // 2,
  207. channels[0],
  208. 3,
  209. stride=1,
  210. padding=1,
  211. act=act)))
  212. else:
  213. self.stem = nn.Sequential(
  214. ('conv1', ConvBNLayer(
  215. 3, channels[0] // 2, 3, stride=2, padding=1, act=act)),
  216. ('conv2', ConvBNLayer(
  217. channels[0] // 2,
  218. channels[0],
  219. 3,
  220. stride=1,
  221. padding=1,
  222. act=act)))
  223. n = len(channels) - 1
  224. self.stages = nn.Sequential(* [(str(i), CSPResStage(
  225. BasicBlock, channels[i], channels[i + 1], layers[i], 2, act=act))
  226. for i in range(n)])
  227. self._out_channels = channels[1:]
  228. self._out_strides = [4, 8, 16, 32]
  229. self.return_idx = return_idx
  230. def forward(self, inputs):
  231. x = inputs['image']
  232. x = self.stem(x)
  233. outs = []
  234. for idx, stage in enumerate(self.stages):
  235. x = stage(x)
  236. if idx in self.return_idx:
  237. outs.append(x)
  238. return outs
  239. @property
  240. def out_shape(self):
  241. return [
  242. ShapeSpec(
  243. channels=self._out_channels[i], stride=self._out_strides[i])
  244. for i in self.return_idx
  245. ]