densenet.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344
  1. # copyright (c) 2021 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 paddle.nn import Conv2D, BatchNorm, Linear, Dropout
  22. from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D
  23. from paddle.nn.initializer import Uniform
  24. import math
  25. from paddlex.ppcls.utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url
  26. MODEL_URLS = {
  27. "DenseNet121":
  28. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet121_pretrained.pdparams",
  29. "DenseNet161":
  30. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet161_pretrained.pdparams",
  31. "DenseNet169":
  32. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet169_pretrained.pdparams",
  33. "DenseNet201":
  34. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet201_pretrained.pdparams",
  35. "DenseNet264":
  36. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet264_pretrained.pdparams",
  37. }
  38. __all__ = list(MODEL_URLS.keys())
  39. class BNACConvLayer(nn.Layer):
  40. def __init__(self,
  41. num_channels,
  42. num_filters,
  43. filter_size,
  44. stride=1,
  45. pad=0,
  46. groups=1,
  47. act="relu",
  48. name=None):
  49. super(BNACConvLayer, self).__init__()
  50. self._batch_norm = BatchNorm(
  51. num_channels,
  52. act=act,
  53. param_attr=ParamAttr(name=name + '_bn_scale'),
  54. bias_attr=ParamAttr(name + '_bn_offset'),
  55. moving_mean_name=name + '_bn_mean',
  56. moving_variance_name=name + '_bn_variance')
  57. self._conv = Conv2D(
  58. in_channels=num_channels,
  59. out_channels=num_filters,
  60. kernel_size=filter_size,
  61. stride=stride,
  62. padding=pad,
  63. groups=groups,
  64. weight_attr=ParamAttr(name=name + "_weights"),
  65. bias_attr=False)
  66. def forward(self, input):
  67. y = self._batch_norm(input)
  68. y = self._conv(y)
  69. return y
  70. class DenseLayer(nn.Layer):
  71. def __init__(self, num_channels, growth_rate, bn_size, dropout, name=None):
  72. super(DenseLayer, self).__init__()
  73. self.dropout = dropout
  74. self.bn_ac_func1 = BNACConvLayer(
  75. num_channels=num_channels,
  76. num_filters=bn_size * growth_rate,
  77. filter_size=1,
  78. pad=0,
  79. stride=1,
  80. name=name + "_x1")
  81. self.bn_ac_func2 = BNACConvLayer(
  82. num_channels=bn_size * growth_rate,
  83. num_filters=growth_rate,
  84. filter_size=3,
  85. pad=1,
  86. stride=1,
  87. name=name + "_x2")
  88. if dropout:
  89. self.dropout_func = Dropout(p=dropout, mode="downscale_in_infer")
  90. def forward(self, input):
  91. conv = self.bn_ac_func1(input)
  92. conv = self.bn_ac_func2(conv)
  93. if self.dropout:
  94. conv = self.dropout_func(conv)
  95. conv = paddle.concat([input, conv], axis=1)
  96. return conv
  97. class DenseBlock(nn.Layer):
  98. def __init__(self,
  99. num_channels,
  100. num_layers,
  101. bn_size,
  102. growth_rate,
  103. dropout,
  104. name=None):
  105. super(DenseBlock, self).__init__()
  106. self.dropout = dropout
  107. self.dense_layer_func = []
  108. pre_channel = num_channels
  109. for layer in range(num_layers):
  110. self.dense_layer_func.append(
  111. self.add_sublayer(
  112. "{}_{}".format(name, layer + 1),
  113. DenseLayer(
  114. num_channels=pre_channel,
  115. growth_rate=growth_rate,
  116. bn_size=bn_size,
  117. dropout=dropout,
  118. name=name + '_' + str(layer + 1))))
  119. pre_channel = pre_channel + growth_rate
  120. def forward(self, input):
  121. conv = input
  122. for func in self.dense_layer_func:
  123. conv = func(conv)
  124. return conv
  125. class TransitionLayer(nn.Layer):
  126. def __init__(self, num_channels, num_output_features, name=None):
  127. super(TransitionLayer, self).__init__()
  128. self.conv_ac_func = BNACConvLayer(
  129. num_channels=num_channels,
  130. num_filters=num_output_features,
  131. filter_size=1,
  132. pad=0,
  133. stride=1,
  134. name=name)
  135. self.pool2d_avg = AvgPool2D(kernel_size=2, stride=2, padding=0)
  136. def forward(self, input):
  137. y = self.conv_ac_func(input)
  138. y = self.pool2d_avg(y)
  139. return y
  140. class ConvBNLayer(nn.Layer):
  141. def __init__(self,
  142. num_channels,
  143. num_filters,
  144. filter_size,
  145. stride=1,
  146. pad=0,
  147. groups=1,
  148. act="relu",
  149. name=None):
  150. super(ConvBNLayer, self).__init__()
  151. self._conv = Conv2D(
  152. in_channels=num_channels,
  153. out_channels=num_filters,
  154. kernel_size=filter_size,
  155. stride=stride,
  156. padding=pad,
  157. groups=groups,
  158. weight_attr=ParamAttr(name=name + "_weights"),
  159. bias_attr=False)
  160. self._batch_norm = BatchNorm(
  161. num_filters,
  162. act=act,
  163. param_attr=ParamAttr(name=name + '_bn_scale'),
  164. bias_attr=ParamAttr(name + '_bn_offset'),
  165. moving_mean_name=name + '_bn_mean',
  166. moving_variance_name=name + '_bn_variance')
  167. def forward(self, input):
  168. y = self._conv(input)
  169. y = self._batch_norm(y)
  170. return y
  171. class DenseNet(nn.Layer):
  172. def __init__(self, layers=60, bn_size=4, dropout=0, class_num=1000):
  173. super(DenseNet, self).__init__()
  174. supported_layers = [121, 161, 169, 201, 264]
  175. assert layers in supported_layers, \
  176. "supported layers are {} but input layer is {}".format(
  177. supported_layers, layers)
  178. densenet_spec = {
  179. 121: (64, 32, [6, 12, 24, 16]),
  180. 161: (96, 48, [6, 12, 36, 24]),
  181. 169: (64, 32, [6, 12, 32, 32]),
  182. 201: (64, 32, [6, 12, 48, 32]),
  183. 264: (64, 32, [6, 12, 64, 48])
  184. }
  185. num_init_features, growth_rate, block_config = densenet_spec[layers]
  186. self.conv1_func = ConvBNLayer(
  187. num_channels=3,
  188. num_filters=num_init_features,
  189. filter_size=7,
  190. stride=2,
  191. pad=3,
  192. act='relu',
  193. name="conv1")
  194. self.pool2d_max = MaxPool2D(kernel_size=3, stride=2, padding=1)
  195. self.block_config = block_config
  196. self.dense_block_func_list = []
  197. self.transition_func_list = []
  198. pre_num_channels = num_init_features
  199. num_features = num_init_features
  200. for i, num_layers in enumerate(block_config):
  201. self.dense_block_func_list.append(
  202. self.add_sublayer(
  203. "db_conv_{}".format(i + 2),
  204. DenseBlock(
  205. num_channels=pre_num_channels,
  206. num_layers=num_layers,
  207. bn_size=bn_size,
  208. growth_rate=growth_rate,
  209. dropout=dropout,
  210. name='conv' + str(i + 2))))
  211. num_features = num_features + num_layers * growth_rate
  212. pre_num_channels = num_features
  213. if i != len(block_config) - 1:
  214. self.transition_func_list.append(
  215. self.add_sublayer(
  216. "tr_conv{}_blk".format(i + 2),
  217. TransitionLayer(
  218. num_channels=pre_num_channels,
  219. num_output_features=num_features // 2,
  220. name='conv' + str(i + 2) + "_blk")))
  221. pre_num_channels = num_features // 2
  222. num_features = num_features // 2
  223. self.batch_norm = BatchNorm(
  224. num_features,
  225. act="relu",
  226. param_attr=ParamAttr(name='conv5_blk_bn_scale'),
  227. bias_attr=ParamAttr(name='conv5_blk_bn_offset'),
  228. moving_mean_name='conv5_blk_bn_mean',
  229. moving_variance_name='conv5_blk_bn_variance')
  230. self.pool2d_avg = AdaptiveAvgPool2D(1)
  231. stdv = 1.0 / math.sqrt(num_features * 1.0)
  232. self.out = Linear(
  233. num_features,
  234. class_num,
  235. weight_attr=ParamAttr(
  236. initializer=Uniform(-stdv, stdv), name="fc_weights"),
  237. bias_attr=ParamAttr(name="fc_offset"))
  238. def forward(self, input):
  239. conv = self.conv1_func(input)
  240. conv = self.pool2d_max(conv)
  241. for i, num_layers in enumerate(self.block_config):
  242. conv = self.dense_block_func_list[i](conv)
  243. if i != len(self.block_config) - 1:
  244. conv = self.transition_func_list[i](conv)
  245. conv = self.batch_norm(conv)
  246. y = self.pool2d_avg(conv)
  247. y = paddle.flatten(y, start_axis=1, stop_axis=-1)
  248. y = self.out(y)
  249. return y
  250. def _load_pretrained(pretrained, model, model_url, use_ssld=False):
  251. if pretrained is False:
  252. pass
  253. elif pretrained is True:
  254. load_dygraph_pretrain_from_url(model, model_url, use_ssld=use_ssld)
  255. elif isinstance(pretrained, str):
  256. load_dygraph_pretrain(model, pretrained)
  257. else:
  258. raise RuntimeError(
  259. "pretrained type is not available. Please use `string` or `boolean` type."
  260. )
  261. def DenseNet121(pretrained=False, use_ssld=False, **kwargs):
  262. model = DenseNet(layers=121, **kwargs)
  263. _load_pretrained(
  264. pretrained, model, MODEL_URLS["DenseNet121"], use_ssld=use_ssld)
  265. return model
  266. def DenseNet161(pretrained=False, use_ssld=False, **kwargs):
  267. model = DenseNet(layers=161, **kwargs)
  268. _load_pretrained(
  269. pretrained, model, MODEL_URLS["DenseNet161"], use_ssld=use_ssld)
  270. return model
  271. def DenseNet169(pretrained=False, use_ssld=False, **kwargs):
  272. model = DenseNet(layers=169, **kwargs)
  273. _load_pretrained(
  274. pretrained, model, MODEL_URLS["DenseNet169"], use_ssld=use_ssld)
  275. return model
  276. def DenseNet201(pretrained=False, use_ssld=False, **kwargs):
  277. model = DenseNet(layers=201, **kwargs)
  278. _load_pretrained(
  279. pretrained, model, MODEL_URLS["DenseNet201"], use_ssld=use_ssld)
  280. return model
  281. def DenseNet264(pretrained=False, use_ssld=False, **kwargs):
  282. model = DenseNet(layers=264, **kwargs)
  283. _load_pretrained(
  284. pretrained, model, MODEL_URLS["DenseNet264"], use_ssld=use_ssld)
  285. return model