Browse Source

update ema

will-jl944 4 years ago
parent
commit
2a076ffec8
2 changed files with 2 additions and 51 deletions
  1. 2 3
      paddlex/cv/models/detector.py
  2. 0 48
      paddlex/cv/models/utils/ema.py

+ 2 - 3
paddlex/cv/models/detector.py

@@ -31,7 +31,7 @@ from paddlex.cv.transforms.batch_operators import BatchCompose, BatchRandomResiz
 from paddlex.cv.transforms import arrange_transforms
 from paddlex.cv.transforms import arrange_transforms
 from .base import BaseModel
 from .base import BaseModel
 from .utils.det_metrics import VOCMetric, COCOMetric
 from .utils.det_metrics import VOCMetric, COCOMetric
-from .utils.ema import ExponentialMovingAverage
+from paddlex.ppdet.optimizer import ModelEMA
 from paddlex.utils.checkpoint import det_pretrain_weights_dict
 from paddlex.utils.checkpoint import det_pretrain_weights_dict
 
 
 __all__ = [
 __all__ = [
@@ -275,8 +275,7 @@ class BaseDetector(BaseModel):
                                  'ESNet_' in self.backbone_name))
                                  'ESNet_' in self.backbone_name))
 
 
         if use_ema:
         if use_ema:
-            ema = ExponentialMovingAverage(
-                decay=.9998, model=self.net, use_thres_step=True)
+            ema = ModelEMA(model=self.net, decay=.9998, use_thres_step=True)
         else:
         else:
             ema = None
             ema = None
         # start train loop
         # start train loop

+ 0 - 48
paddlex/cv/models/utils/ema.py

@@ -1,48 +0,0 @@
-# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import paddle
-
-
-class ExponentialMovingAverage(object):
-    def __init__(self, decay, model, use_thres_step=False):
-        self.step = 0
-        self.decay = decay
-        self.shadow = dict()
-        for k, v in model.state_dict().items():
-            self.shadow[k] = paddle.zeros_like(v)
-        self.use_thres_step = use_thres_step
-
-    def update(self, model):
-        if self.use_thres_step:
-            decay = min(self.decay, (1 + self.step) / (10 + self.step))
-        else:
-            decay = self.decay
-        self._decay = decay
-        model_dict = model.state_dict()
-        for k, v in self.shadow.items():
-            v = decay * v + (1 - decay) * model_dict[k]
-            v.stop_gradient = True
-            self.shadow[k] = v
-        self.step += 1
-
-    def apply(self):
-        if self.step == 0:
-            return self.shadow
-        state_dict = dict()
-        for k, v in self.shadow.items():
-            v = v / (1 - self._decay**self.step)
-            v.stop_gradient = True
-            state_dict[k] = v
-        return state_dict