| | |
| | import math |
| | from typing import Optional |
| |
|
| | import torch |
| | import torch.nn as nn |
| | from mmengine.model import ExponentialMovingAverage |
| | from torch import Tensor |
| |
|
| | from mmdet.registry import MODELS |
| |
|
| |
|
| | @MODELS.register_module() |
| | class ExpMomentumEMA(ExponentialMovingAverage): |
| | """Exponential moving average (EMA) with exponential momentum strategy, |
| | which is used in YOLOX. |
| | |
| | Args: |
| | model (nn.Module): The model to be averaged. |
| | momentum (float): The momentum used for updating ema parameter. |
| | Ema's parameter are updated with the formula: |
| | `averaged_param = (1-momentum) * averaged_param + momentum * |
| | source_param`. Defaults to 0.0002. |
| | gamma (int): Use a larger momentum early in training and gradually |
| | annealing to a smaller value to update the ema model smoothly. The |
| | momentum is calculated as |
| | `(1 - momentum) * exp(-(1 + steps) / gamma) + momentum`. |
| | Defaults to 2000. |
| | interval (int): Interval between two updates. Defaults to 1. |
| | device (torch.device, optional): If provided, the averaged model will |
| | be stored on the :attr:`device`. Defaults to None. |
| | update_buffers (bool): if True, it will compute running averages for |
| | both the parameters and the buffers of the model. Defaults to |
| | False. |
| | """ |
| |
|
| | def __init__(self, |
| | model: nn.Module, |
| | momentum: float = 0.0002, |
| | gamma: int = 2000, |
| | interval=1, |
| | device: Optional[torch.device] = None, |
| | update_buffers: bool = False) -> None: |
| | super().__init__( |
| | model=model, |
| | momentum=momentum, |
| | interval=interval, |
| | device=device, |
| | update_buffers=update_buffers) |
| | assert gamma > 0, f'gamma must be greater than 0, but got {gamma}' |
| | self.gamma = gamma |
| |
|
| | def avg_func(self, averaged_param: Tensor, source_param: Tensor, |
| | steps: int) -> None: |
| | """Compute the moving average of the parameters using the exponential |
| | momentum strategy. |
| | |
| | Args: |
| | averaged_param (Tensor): The averaged parameters. |
| | source_param (Tensor): The source parameters. |
| | steps (int): The number of times the parameters have been |
| | updated. |
| | """ |
| | momentum = (1 - self.momentum) * math.exp( |
| | -float(1 + steps) / self.gamma) + self.momentum |
| | averaged_param.lerp_(source_param, momentum) |
| |
|
