initializer.py 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307
  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 math
  15. import numpy as np
  16. import paddle
  17. import paddle.nn as nn
  18. __all__ = [
  19. 'uniform_',
  20. 'normal_',
  21. 'constant_',
  22. 'ones_',
  23. 'zeros_',
  24. 'xavier_uniform_',
  25. 'xavier_normal_',
  26. 'kaiming_uniform_',
  27. 'kaiming_normal_',
  28. 'linear_init_',
  29. 'conv_init_',
  30. 'reset_initialized_parameter',
  31. ]
  32. def _no_grad_uniform_(tensor, a, b):
  33. with paddle.no_grad():
  34. tensor.set_value(
  35. paddle.uniform(
  36. shape=tensor.shape, dtype=tensor.dtype, min=a, max=b))
  37. return tensor
  38. def _no_grad_normal_(tensor, mean=0., std=1.):
  39. with paddle.no_grad():
  40. tensor.set_value(paddle.normal(mean=mean, std=std, shape=tensor.shape))
  41. return tensor
  42. def _no_grad_fill_(tensor, value=0.):
  43. with paddle.no_grad():
  44. tensor.set_value(paddle.full_like(tensor, value, dtype=tensor.dtype))
  45. return tensor
  46. def uniform_(tensor, a, b):
  47. """
  48. Modified tensor inspace using uniform_
  49. Args:
  50. tensor (paddle.Tensor): paddle Tensor
  51. a (float|int): min value.
  52. b (float|int): max value.
  53. Return:
  54. tensor
  55. """
  56. return _no_grad_uniform_(tensor, a, b)
  57. def normal_(tensor, mean=0., std=1.):
  58. """
  59. Modified tensor inspace using normal_
  60. Args:
  61. tensor (paddle.Tensor): paddle Tensor
  62. mean (float|int): mean value.
  63. std (float|int): std value.
  64. Return:
  65. tensor
  66. """
  67. return _no_grad_normal_(tensor, mean, std)
  68. def constant_(tensor, value=0.):
  69. """
  70. Modified tensor inspace using constant_
  71. Args:
  72. tensor (paddle.Tensor): paddle Tensor
  73. value (float|int): value to fill tensor.
  74. Return:
  75. tensor
  76. """
  77. return _no_grad_fill_(tensor, value)
  78. def ones_(tensor):
  79. """
  80. Modified tensor inspace using ones_
  81. Args:
  82. tensor (paddle.Tensor): paddle Tensor
  83. Return:
  84. tensor
  85. """
  86. return _no_grad_fill_(tensor, 1)
  87. def zeros_(tensor):
  88. """
  89. Modified tensor inspace using zeros_
  90. Args:
  91. tensor (paddle.Tensor): paddle Tensor
  92. Return:
  93. tensor
  94. """
  95. return _no_grad_fill_(tensor, 0)
  96. def _calculate_fan_in_and_fan_out(tensor, reverse=False):
  97. """
  98. Calculate (fan_in, _fan_out) for tensor
  99. Args:
  100. tensor (Tensor): paddle.Tensor
  101. reverse (bool: False): tensor data format order, False by default as [fout, fin, ...]. e.g. : conv.weight [cout, cin, kh, kw] is False; linear.weight [cin, cout] is True
  102. Return:
  103. Tuple[fan_in, fan_out]
  104. """
  105. if tensor.ndim < 2:
  106. raise ValueError(
  107. "Fan in and fan out can not be computed for tensor with fewer than 2 dimensions"
  108. )
  109. if reverse:
  110. num_input_fmaps, num_output_fmaps = tensor.shape[0], tensor.shape[1]
  111. else:
  112. num_input_fmaps, num_output_fmaps = tensor.shape[1], tensor.shape[0]
  113. receptive_field_size = 1
  114. if tensor.ndim > 2:
  115. receptive_field_size = np.prod(tensor.shape[2:])
  116. fan_in = num_input_fmaps * receptive_field_size
  117. fan_out = num_output_fmaps * receptive_field_size
  118. return fan_in, fan_out
  119. def xavier_uniform_(tensor, gain=1., reverse=False):
  120. """
  121. Modified tensor inspace using xavier_uniform_
  122. Args:
  123. tensor (paddle.Tensor): paddle Tensor
  124. gain (float): super parameter, 1. default.
  125. reverse (bool): reverse (bool: False): tensor data format order, False by default as [fout, fin, ...].
  126. Return:
  127. tensor
  128. """
  129. fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse=reverse)
  130. std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
  131. k = math.sqrt(3.0) * std
  132. return _no_grad_uniform_(tensor, -k, k)
  133. def xavier_normal_(tensor, gain=1., reverse=False):
  134. """
  135. Modified tensor inspace using xavier_normal_
  136. Args:
  137. tensor (paddle.Tensor): paddle Tensor
  138. gain (float): super parameter, 1. default.
  139. reverse (bool): reverse (bool: False): tensor data format order, False by default as [fout, fin, ...].
  140. Return:
  141. tensor
  142. """
  143. fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse=reverse)
  144. std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
  145. return _no_grad_normal_(tensor, 0, std)
  146. # reference: https://pytorch.org/docs/stable/_modules/torch/nn/init.html
  147. def _calculate_correct_fan(tensor, mode, reverse=False):
  148. mode = mode.lower()
  149. valid_modes = ['fan_in', 'fan_out']
  150. if mode not in valid_modes:
  151. raise ValueError("Mode {} not supported, please use one of {}".format(
  152. mode, valid_modes))
  153. fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse)
  154. return fan_in if mode == 'fan_in' else fan_out
  155. def _calculate_gain(nonlinearity, param=None):
  156. linear_fns = [
  157. 'linear', 'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d',
  158. 'conv_transpose2d', 'conv_transpose3d'
  159. ]
  160. if nonlinearity in linear_fns or nonlinearity == 'sigmoid':
  161. return 1
  162. elif nonlinearity == 'tanh':
  163. return 5.0 / 3
  164. elif nonlinearity == 'relu':
  165. return math.sqrt(2.0)
  166. elif nonlinearity == 'leaky_relu':
  167. if param is None:
  168. negative_slope = 0.01
  169. elif not isinstance(param, bool) and isinstance(
  170. param, int) or isinstance(param, float):
  171. # True/False are instances of int, hence check above
  172. negative_slope = param
  173. else:
  174. raise ValueError("negative_slope {} not a valid number".format(
  175. param))
  176. return math.sqrt(2.0 / (1 + negative_slope**2))
  177. elif nonlinearity == 'selu':
  178. return 3.0 / 4
  179. else:
  180. raise ValueError("Unsupported nonlinearity {}".format(nonlinearity))
  181. def kaiming_uniform_(tensor,
  182. a=0,
  183. mode='fan_in',
  184. nonlinearity='leaky_relu',
  185. reverse=False):
  186. """
  187. Modified tensor inspace using kaiming_uniform method
  188. Args:
  189. tensor (paddle.Tensor): paddle Tensor
  190. mode (str): ['fan_in', 'fan_out'], 'fin_in' defalut
  191. nonlinearity (str): nonlinearity method name
  192. reverse (bool): reverse (bool: False): tensor data format order, False by default as [fout, fin, ...].
  193. Return:
  194. tensor
  195. """
  196. fan = _calculate_correct_fan(tensor, mode, reverse)
  197. gain = _calculate_gain(nonlinearity, a)
  198. std = gain / math.sqrt(fan)
  199. k = math.sqrt(3.0) * std
  200. return _no_grad_uniform_(tensor, -k, k)
  201. def kaiming_normal_(tensor,
  202. a=0,
  203. mode='fan_in',
  204. nonlinearity='leaky_relu',
  205. reverse=False):
  206. """
  207. Modified tensor inspace using kaiming_normal_
  208. Args:
  209. tensor (paddle.Tensor): paddle Tensor
  210. mode (str): ['fan_in', 'fan_out'], 'fin_in' defalut
  211. nonlinearity (str): nonlinearity method name
  212. reverse (bool): reverse (bool: False): tensor data format order, False by default as [fout, fin, ...].
  213. Return:
  214. tensor
  215. """
  216. fan = _calculate_correct_fan(tensor, mode, reverse)
  217. gain = _calculate_gain(nonlinearity, a)
  218. std = gain / math.sqrt(fan)
  219. return _no_grad_normal_(tensor, 0, std)
  220. def linear_init_(module):
  221. bound = 1 / math.sqrt(module.weight.shape[0])
  222. uniform_(module.weight, -bound, bound)
  223. uniform_(module.bias, -bound, bound)
  224. def conv_init_(module):
  225. bound = 1 / np.sqrt(np.prod(module.weight.shape[1:]))
  226. uniform_(module.weight, -bound, bound)
  227. uniform_(module.bias, -bound, bound)
  228. @paddle.no_grad()
  229. def reset_initialized_parameter(model, include_self=True):
  230. """
  231. Reset initialized parameter using following method for [conv, linear, embedding, bn]
  232. Args:
  233. model (paddle.Layer): paddle Layer
  234. include_self (bool: False): include_self for Layer.named_sublayers method. Indicate whether including itself
  235. Return:
  236. None
  237. """
  238. for _, m in model.named_sublayers(include_self=include_self):
  239. if isinstance(m, nn.Conv2D):
  240. k = float(m._groups) / (m._in_channels * m._kernel_size[0] *
  241. m._kernel_size[1])
  242. k = math.sqrt(k)
  243. _no_grad_uniform_(m.weight, -k, k)
  244. if hasattr(m, 'bias') and getattr(m, 'bias') is not None:
  245. _no_grad_uniform_(m.bias, -k, k)
  246. elif isinstance(m, nn.Linear):
  247. k = math.sqrt(1. / m.weight.shape[0])
  248. _no_grad_uniform_(m.weight, -k, k)
  249. if hasattr(m, 'bias') and getattr(m, 'bias') is not None:
  250. _no_grad_uniform_(m.bias, -k, k)
  251. elif isinstance(m, nn.Embedding):
  252. _no_grad_normal_(m.weight, mean=0., std=1.)
  253. elif isinstance(m, (nn.BatchNorm2D, nn.LayerNorm)):
  254. _no_grad_fill_(m.weight, 1.)
  255. if hasattr(m, 'bias') and getattr(m, 'bias') is not None:
  256. _no_grad_fill_(m.bias, 0)