ema.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. # Copyright (c) 2022 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 math
  18. import paddle
  19. import weakref
  20. class ModelEMA(object):
  21. """
  22. Exponential Weighted Average for Deep Neutal Networks
  23. Args:
  24. model (nn.Layer): Detector of model.
  25. decay (int): The decay used for updating ema parameter.
  26. Ema's parameter are updated with the formula:
  27. `ema_param = decay * ema_param + (1 - decay) * cur_param`.
  28. Defaults is 0.9998.
  29. ema_decay_type (str): type in ['threshold', 'normal', 'exponential'],
  30. 'threshold' as default.
  31. cycle_epoch (int): The epoch of interval to reset ema_param and
  32. step. Defaults is -1, which means not reset. Its function is to
  33. add a regular effect to ema, which is set according to experience
  34. and is effective when the total training epoch is large.
  35. """
  36. def __init__(self,
  37. model,
  38. decay=0.9998,
  39. ema_decay_type='threshold',
  40. cycle_epoch=-1):
  41. self.step = 0
  42. self.epoch = 0
  43. self.decay = decay
  44. self.state_dict = dict()
  45. for k, v in model.state_dict().items():
  46. self.state_dict[k] = paddle.zeros_like(v)
  47. self.ema_decay_type = ema_decay_type
  48. self.cycle_epoch = cycle_epoch
  49. self._model_state = {
  50. k: weakref.ref(p)
  51. for k, p in model.state_dict().items()
  52. }
  53. def reset(self):
  54. self.step = 0
  55. self.epoch = 0
  56. for k, v in self.state_dict.items():
  57. self.state_dict[k] = paddle.zeros_like(v)
  58. def resume(self, state_dict, step=0):
  59. for k, v in state_dict.items():
  60. if k in self.state_dict:
  61. self.state_dict[k] = v
  62. self.step = step
  63. def update(self, model=None):
  64. if self.ema_decay_type == 'threshold':
  65. decay = min(self.decay, (1 + self.step) / (10 + self.step))
  66. elif self.ema_decay_type == 'exponential':
  67. decay = self.decay * (1 - math.exp(-(self.step + 1) / 2000))
  68. else:
  69. decay = self.decay
  70. self._decay = decay
  71. if model is not None:
  72. model_dict = model.state_dict()
  73. else:
  74. model_dict = {k: p() for k, p in self._model_state.items()}
  75. assert all(
  76. [v is not None for _, v in model_dict.items()]), 'python gc.'
  77. for k, v in self.state_dict.items():
  78. v = decay * v + (1 - decay) * model_dict[k]
  79. v.stop_gradient = True
  80. self.state_dict[k] = v
  81. self.step += 1
  82. def apply(self):
  83. if self.step == 0:
  84. return self.state_dict
  85. state_dict = dict()
  86. for k, v in self.state_dict.items():
  87. if self.ema_decay_type != 'exponential':
  88. v = v / (1 - self._decay**self.step)
  89. v.stop_gradient = True
  90. state_dict[k] = v
  91. self.epoch += 1
  92. if self.cycle_epoch > 0 and self.epoch == self.cycle_epoch:
  93. self.reset()
  94. return state_dict