rexnet.py 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281
  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. from paddlex.ppcls.utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url
  23. MODEL_URLS = {
  24. "ReXNet_1_0":
  25. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/ReXNet_1_0_pretrained.pdparams",
  26. "ReXNet_1_3":
  27. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/ReXNet_1_3_pretrained.pdparams",
  28. "ReXNet_1_5":
  29. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/ReXNet_1_5_pretrained.pdparams",
  30. "ReXNet_2_0":
  31. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/ReXNet_2_0_pretrained.pdparams",
  32. "ReXNet_3_0":
  33. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/ReXNet_3_0_pretrained.pdparams",
  34. }
  35. __all__ = list(MODEL_URLS.keys())
  36. def conv_bn_act(out,
  37. in_channels,
  38. channels,
  39. kernel=1,
  40. stride=1,
  41. pad=0,
  42. num_group=1,
  43. active=True,
  44. relu6=False):
  45. out.append(
  46. nn.Conv2D(
  47. in_channels,
  48. channels,
  49. kernel,
  50. stride,
  51. pad,
  52. groups=num_group,
  53. bias_attr=False))
  54. out.append(nn.BatchNorm2D(channels))
  55. if active:
  56. out.append(nn.ReLU6() if relu6 else nn.ReLU())
  57. def conv_bn_swish(out,
  58. in_channels,
  59. channels,
  60. kernel=1,
  61. stride=1,
  62. pad=0,
  63. num_group=1):
  64. out.append(
  65. nn.Conv2D(
  66. in_channels,
  67. channels,
  68. kernel,
  69. stride,
  70. pad,
  71. groups=num_group,
  72. bias_attr=False))
  73. out.append(nn.BatchNorm2D(channels))
  74. out.append(nn.Swish())
  75. class SE(nn.Layer):
  76. def __init__(self, in_channels, channels, se_ratio=12):
  77. super(SE, self).__init__()
  78. self.avg_pool = nn.AdaptiveAvgPool2D(1)
  79. self.fc = nn.Sequential(
  80. nn.Conv2D(
  81. in_channels, channels // se_ratio, kernel_size=1, padding=0),
  82. nn.BatchNorm2D(channels // se_ratio),
  83. nn.ReLU(),
  84. nn.Conv2D(
  85. channels // se_ratio, channels, kernel_size=1, padding=0),
  86. nn.Sigmoid())
  87. def forward(self, x):
  88. y = self.avg_pool(x)
  89. y = self.fc(y)
  90. return x * y
  91. class LinearBottleneck(nn.Layer):
  92. def __init__(self,
  93. in_channels,
  94. channels,
  95. t,
  96. stride,
  97. use_se=True,
  98. se_ratio=12,
  99. **kwargs):
  100. super(LinearBottleneck, self).__init__(**kwargs)
  101. self.use_shortcut = stride == 1 and in_channels <= channels
  102. self.in_channels = in_channels
  103. self.out_channels = channels
  104. out = []
  105. if t != 1:
  106. dw_channels = in_channels * t
  107. conv_bn_swish(out, in_channels=in_channels, channels=dw_channels)
  108. else:
  109. dw_channels = in_channels
  110. conv_bn_act(
  111. out,
  112. in_channels=dw_channels,
  113. channels=dw_channels,
  114. kernel=3,
  115. stride=stride,
  116. pad=1,
  117. num_group=dw_channels,
  118. active=False)
  119. if use_se:
  120. out.append(SE(dw_channels, dw_channels, se_ratio))
  121. out.append(nn.ReLU6())
  122. conv_bn_act(
  123. out,
  124. in_channels=dw_channels,
  125. channels=channels,
  126. active=False,
  127. relu6=True)
  128. self.out = nn.Sequential(*out)
  129. def forward(self, x):
  130. out = self.out(x)
  131. if self.use_shortcut:
  132. out[:, 0:self.in_channels] += x
  133. return out
  134. class ReXNetV1(nn.Layer):
  135. def __init__(self,
  136. input_ch=16,
  137. final_ch=180,
  138. width_mult=1.0,
  139. depth_mult=1.0,
  140. class_num=1000,
  141. use_se=True,
  142. se_ratio=12,
  143. dropout_ratio=0.2,
  144. bn_momentum=0.9):
  145. super(ReXNetV1, self).__init__()
  146. layers = [1, 2, 2, 3, 3, 5]
  147. strides = [1, 2, 2, 2, 1, 2]
  148. use_ses = [False, False, True, True, True, True]
  149. layers = [ceil(element * depth_mult) for element in layers]
  150. strides = sum([[element] + [1] * (layers[idx] - 1)
  151. for idx, element in enumerate(strides)], [])
  152. if use_se:
  153. use_ses = sum([[element] * layers[idx]
  154. for idx, element in enumerate(use_ses)], [])
  155. else:
  156. use_ses = [False] * sum(layers[:])
  157. ts = [1] * layers[0] + [6] * sum(layers[1:])
  158. self.depth = sum(layers[:]) * 3
  159. stem_channel = 32 / width_mult if width_mult < 1.0 else 32
  160. inplanes = input_ch / width_mult if width_mult < 1.0 else input_ch
  161. features = []
  162. in_channels_group = []
  163. channels_group = []
  164. # The following channel configuration is a simple instance to make each layer become an expand layer.
  165. for i in range(self.depth // 3):
  166. if i == 0:
  167. in_channels_group.append(int(round(stem_channel * width_mult)))
  168. channels_group.append(int(round(inplanes * width_mult)))
  169. else:
  170. in_channels_group.append(int(round(inplanes * width_mult)))
  171. inplanes += final_ch / (self.depth // 3 * 1.0)
  172. channels_group.append(int(round(inplanes * width_mult)))
  173. conv_bn_swish(
  174. features,
  175. 3,
  176. int(round(stem_channel * width_mult)),
  177. kernel=3,
  178. stride=2,
  179. pad=1)
  180. for block_idx, (in_c, c, t, s, se) in enumerate(
  181. zip(in_channels_group, channels_group, ts, strides, use_ses)):
  182. features.append(
  183. LinearBottleneck(
  184. in_channels=in_c,
  185. channels=c,
  186. t=t,
  187. stride=s,
  188. use_se=se,
  189. se_ratio=se_ratio))
  190. pen_channels = int(1280 * width_mult)
  191. conv_bn_swish(features, c, pen_channels)
  192. features.append(nn.AdaptiveAvgPool2D(1))
  193. self.features = nn.Sequential(*features)
  194. self.output = nn.Sequential(
  195. nn.Dropout(dropout_ratio),
  196. nn.Conv2D(
  197. pen_channels, class_num, 1, bias_attr=True))
  198. def forward(self, x):
  199. x = self.features(x)
  200. x = self.output(x).squeeze(axis=-1).squeeze(axis=-1)
  201. return x
  202. def _load_pretrained(pretrained, model, model_url, use_ssld=False):
  203. if pretrained is False:
  204. pass
  205. elif pretrained is True:
  206. load_dygraph_pretrain_from_url(model, model_url, use_ssld=use_ssld)
  207. elif isinstance(pretrained, str):
  208. load_dygraph_pretrain(model, pretrained)
  209. else:
  210. raise RuntimeError(
  211. "pretrained type is not available. Please use `string` or `boolean` type."
  212. )
  213. def ReXNet_1_0(pretrained=False, use_ssld=False, **kwargs):
  214. model = ReXNetV1(width_mult=1.0, **kwargs)
  215. _load_pretrained(
  216. pretrained, model, MODEL_URLS["ReXNet_1_0"], use_ssld=use_ssld)
  217. return model
  218. def ReXNet_1_3(pretrained=False, use_ssld=False, **kwargs):
  219. model = ReXNetV1(width_mult=1.3, **kwargs)
  220. _load_pretrained(
  221. pretrained, model, MODEL_URLS["ReXNet_1_3"], use_ssld=use_ssld)
  222. return model
  223. def ReXNet_1_5(pretrained=False, use_ssld=False, **kwargs):
  224. model = ReXNetV1(width_mult=1.5, **kwargs)
  225. _load_pretrained(
  226. pretrained, model, MODEL_URLS["ReXNet_1_5"], use_ssld=use_ssld)
  227. return model
  228. def ReXNet_2_0(pretrained=False, use_ssld=False, **kwargs):
  229. model = ReXNetV1(width_mult=2.0, **kwargs)
  230. _load_pretrained(
  231. pretrained, model, MODEL_URLS["ReXNet_2_0"], use_ssld=use_ssld)
  232. return model
  233. def ReXNet_3_0(pretrained=False, use_ssld=False, **kwargs):
  234. model = ReXNetV1(width_mult=3.0, **kwargs)
  235. _load_pretrained(
  236. pretrained, model, MODEL_URLS["ReXNet_3_0"], use_ssld=use_ssld)
  237. return model