rexnet.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240
  1. # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  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 numpy as np
  18. import paddle
  19. from paddle import ParamAttr
  20. import paddle.nn as nn
  21. from math import ceil
  22. __all__ = [
  23. "ReXNet_1_0", "ReXNet_1_3", "ReXNet_1_5", "ReXNet_2_0", "ReXNet_3_0"
  24. ]
  25. def conv_bn_act(out,
  26. in_channels,
  27. channels,
  28. kernel=1,
  29. stride=1,
  30. pad=0,
  31. num_group=1,
  32. active=True,
  33. relu6=False):
  34. out.append(
  35. nn.Conv2D(
  36. in_channels,
  37. channels,
  38. kernel,
  39. stride,
  40. pad,
  41. groups=num_group,
  42. bias_attr=False))
  43. out.append(nn.BatchNorm2D(channels))
  44. if active:
  45. out.append(nn.ReLU6() if relu6 else nn.ReLU())
  46. def conv_bn_swish(out,
  47. in_channels,
  48. channels,
  49. kernel=1,
  50. stride=1,
  51. pad=0,
  52. num_group=1):
  53. out.append(
  54. nn.Conv2D(
  55. in_channels,
  56. channels,
  57. kernel,
  58. stride,
  59. pad,
  60. groups=num_group,
  61. bias_attr=False))
  62. out.append(nn.BatchNorm2D(channels))
  63. out.append(nn.Swish())
  64. class SE(nn.Layer):
  65. def __init__(self, in_channels, channels, se_ratio=12):
  66. super(SE, self).__init__()
  67. self.avg_pool = nn.AdaptiveAvgPool2D(1)
  68. self.fc = nn.Sequential(
  69. nn.Conv2D(
  70. in_channels, channels // se_ratio, kernel_size=1, padding=0),
  71. nn.BatchNorm2D(channels // se_ratio),
  72. nn.ReLU(),
  73. nn.Conv2D(
  74. channels // se_ratio, channels, kernel_size=1, padding=0),
  75. nn.Sigmoid())
  76. def forward(self, x):
  77. y = self.avg_pool(x)
  78. y = self.fc(y)
  79. return x * y
  80. class LinearBottleneck(nn.Layer):
  81. def __init__(self,
  82. in_channels,
  83. channels,
  84. t,
  85. stride,
  86. use_se=True,
  87. se_ratio=12,
  88. **kwargs):
  89. super(LinearBottleneck, self).__init__(**kwargs)
  90. self.use_shortcut = stride == 1 and in_channels <= channels
  91. self.in_channels = in_channels
  92. self.out_channels = channels
  93. out = []
  94. if t != 1:
  95. dw_channels = in_channels * t
  96. conv_bn_swish(out, in_channels=in_channels, channels=dw_channels)
  97. else:
  98. dw_channels = in_channels
  99. conv_bn_act(
  100. out,
  101. in_channels=dw_channels,
  102. channels=dw_channels,
  103. kernel=3,
  104. stride=stride,
  105. pad=1,
  106. num_group=dw_channels,
  107. active=False)
  108. if use_se:
  109. out.append(SE(dw_channels, dw_channels, se_ratio))
  110. out.append(nn.ReLU6())
  111. conv_bn_act(
  112. out,
  113. in_channels=dw_channels,
  114. channels=channels,
  115. active=False,
  116. relu6=True)
  117. self.out = nn.Sequential(*out)
  118. def forward(self, x):
  119. out = self.out(x)
  120. if self.use_shortcut:
  121. out[:, 0:self.in_channels] += x
  122. return out
  123. class ReXNetV1(nn.Layer):
  124. def __init__(self,
  125. input_ch=16,
  126. final_ch=180,
  127. width_mult=1.0,
  128. depth_mult=1.0,
  129. class_dim=1000,
  130. use_se=True,
  131. se_ratio=12,
  132. dropout_ratio=0.2,
  133. bn_momentum=0.9):
  134. super(ReXNetV1, self).__init__()
  135. layers = [1, 2, 2, 3, 3, 5]
  136. strides = [1, 2, 2, 2, 1, 2]
  137. use_ses = [False, False, True, True, True, True]
  138. layers = [ceil(element * depth_mult) for element in layers]
  139. strides = sum([[element] + [1] * (layers[idx] - 1)
  140. for idx, element in enumerate(strides)], [])
  141. if use_se:
  142. use_ses = sum([[element] * layers[idx]
  143. for idx, element in enumerate(use_ses)], [])
  144. else:
  145. use_ses = [False] * sum(layers[:])
  146. ts = [1] * layers[0] + [6] * sum(layers[1:])
  147. self.depth = sum(layers[:]) * 3
  148. stem_channel = 32 / width_mult if width_mult < 1.0 else 32
  149. inplanes = input_ch / width_mult if width_mult < 1.0 else input_ch
  150. features = []
  151. in_channels_group = []
  152. channels_group = []
  153. # The following channel configuration is a simple instance to make each layer become an expand layer.
  154. for i in range(self.depth // 3):
  155. if i == 0:
  156. in_channels_group.append(int(round(stem_channel * width_mult)))
  157. channels_group.append(int(round(inplanes * width_mult)))
  158. else:
  159. in_channels_group.append(int(round(inplanes * width_mult)))
  160. inplanes += final_ch / (self.depth // 3 * 1.0)
  161. channels_group.append(int(round(inplanes * width_mult)))
  162. conv_bn_swish(
  163. features,
  164. 3,
  165. int(round(stem_channel * width_mult)),
  166. kernel=3,
  167. stride=2,
  168. pad=1)
  169. for block_idx, (in_c, c, t, s, se) in enumerate(
  170. zip(in_channels_group, channels_group, ts, strides, use_ses)):
  171. features.append(
  172. LinearBottleneck(
  173. in_channels=in_c,
  174. channels=c,
  175. t=t,
  176. stride=s,
  177. use_se=se,
  178. se_ratio=se_ratio))
  179. pen_channels = int(1280 * width_mult)
  180. conv_bn_swish(features, c, pen_channels)
  181. features.append(nn.AdaptiveAvgPool2D(1))
  182. self.features = nn.Sequential(*features)
  183. self.output = nn.Sequential(
  184. nn.Dropout(dropout_ratio),
  185. nn.Conv2D(
  186. pen_channels, class_dim, 1, bias_attr=True))
  187. def forward(self, x):
  188. x = self.features(x)
  189. x = self.output(x).squeeze(axis=-1).squeeze(axis=-1)
  190. return x
  191. def ReXNet_1_0(**args):
  192. return ReXNetV1(width_mult=1.0, **args)
  193. def ReXNet_1_3(**args):
  194. return ReXNetV1(width_mult=1.3, **args)
  195. def ReXNet_1_5(**args):
  196. return ReXNetV1(width_mult=1.5, **args)
  197. def ReXNet_2_0(**args):
  198. return ReXNetV1(width_mult=2.0, **args)
  199. def ReXNet_3_0(**args):
  200. return ReXNetV1(width_mult=3.0, **args)