Shortcuts

mmengine.model.averaged_model 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import logging
from abc import abstractmethod
from copy import deepcopy
from typing import Optional

import torch
import torch.nn as nn
from torch import Tensor

from mmengine.logging import print_log
from mmengine.registry import MODELS


[文档]class BaseAveragedModel(nn.Module): """A base class for averaging model weights. Weight averaging, such as SWA and EMA, is a widely used technique for training neural networks. This class implements the averaging process for a model. All subclasses must implement the `avg_func` method. This class creates a copy of the provided module :attr:`model` on the :attr:`device` and allows computing running averages of the parameters of the :attr:`model`. The code is referenced from: https://github.com/pytorch/pytorch/blob/master/torch/optim/swa_utils.py. Different from the `AveragedModel` in PyTorch, we use in-place operation to improve the parameter updating speed, which is about 5 times faster than the non-in-place version. In mmengine, we provide two ways to use the model averaging: 1. Use the model averaging module in hook: We provide an :class:`mmengine.hooks.EMAHook` to apply the model averaging during training. Add ``custom_hooks=[dict(type='EMAHook')]`` to the config or the runner. 2. Use the model averaging module directly in the algorithm. Take the ema teacher in semi-supervise as an example: >>> from mmengine.model import ExponentialMovingAverage >>> student = ResNet(depth=50) >>> # use ema model as teacher >>> ema_teacher = ExponentialMovingAverage(student) Args: model (nn.Module): The model to be averaged. 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. """ # noqa: E501 def __init__(self, model: nn.Module, interval: int = 1, device: Optional[torch.device] = None, update_buffers: bool = False) -> None: super().__init__() self.module = deepcopy(model).requires_grad_(False) self.interval = interval if device is not None: self.module = self.module.to(device) self.register_buffer('steps', torch.tensor(0, dtype=torch.long, device=device)) self.update_buffers = update_buffers if update_buffers: self.avg_parameters = self.module.state_dict() else: self.avg_parameters = dict(self.module.named_parameters())
[文档] @abstractmethod def avg_func(self, averaged_param: Tensor, source_param: Tensor, steps: int) -> None: """Use in-place operation to compute the average of the parameters. All subclasses must implement this method. Args: averaged_param (Tensor): The averaged parameters. source_param (Tensor): The source parameters. steps (int): The number of times the parameters have been updated. """
[文档] def forward(self, *args, **kwargs): """Forward method of the averaged model.""" return self.module(*args, **kwargs)
[文档] def update_parameters(self, model: nn.Module) -> None: """Update the parameters of the model. This method will execute the ``avg_func`` to compute the new parameters and update the model's parameters. Args: model (nn.Module): The model whose parameters will be averaged. """ src_parameters = ( model.state_dict() if self.update_buffers else dict(model.named_parameters())) if self.steps == 0: for k, p_avg in self.avg_parameters.items(): p_avg.data.copy_(src_parameters[k].data) elif self.steps % self.interval == 0: for k, p_avg in self.avg_parameters.items(): if p_avg.dtype.is_floating_point: device = p_avg.device self.avg_func(p_avg.data, src_parameters[k].data.to(device), self.steps) if not self.update_buffers: # If not update the buffers, # keep the buffers in sync with the source model. for b_avg, b_src in zip(self.module.buffers(), model.buffers()): b_avg.data.copy_(b_src.data.to(b_avg.device)) self.steps += 1
[文档]@MODELS.register_module() class StochasticWeightAverage(BaseAveragedModel): """Implements the stochastic weight averaging (SWA) of the model. Stochastic Weight Averaging was proposed in `Averaging Weights Leads to Wider Optima and Better Generalization, UAI 2018. <https://arxiv.org/abs/1803.05407>`_ by Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov and Andrew Gordon Wilson. """
[文档] def avg_func(self, averaged_param: Tensor, source_param: Tensor, steps: int) -> None: """Compute the average of the parameters using stochastic weight average. Args: averaged_param (Tensor): The averaged parameters. source_param (Tensor): The source parameters. steps (int): The number of times the parameters have been updated. """ averaged_param.add_( source_param - averaged_param, alpha=1 / float(steps // self.interval + 1))
[文档]@MODELS.register_module() class ExponentialMovingAverage(BaseAveragedModel): r"""Implements the exponential moving average (EMA) of the model. All parameters are updated by the formula as below: .. math:: Xema_{t+1} = (1 - momentum) * Xema_{t} + momentum * X_t .. note:: This :attr:`momentum` argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, :math:`Xema_{t+1}` is the moving average and :math:`X_t` is the new observed value. The value of momentum is usually a small number, allowing observed values to slowly update the ema parameters. Args: model (nn.Module): The model to be averaged. momentum (float): The momentum used for updating ema parameter. Defaults to 0.0002. Ema's parameter are updated with the formula :math:`averaged\_param = (1-momentum) * averaged\_param + momentum * source\_param`. 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. """ # noqa: W605 def __init__(self, model: nn.Module, momentum: float = 0.0002, interval: int = 1, device: Optional[torch.device] = None, update_buffers: bool = False) -> None: super().__init__(model, interval, device, update_buffers) assert 0.0 < momentum < 1.0, 'momentum must be in range (0.0, 1.0)'\ f'but got {momentum}' if momentum > 0.5: print_log( 'The value of momentum in EMA is usually a small number,' 'which is different from the conventional notion of ' f'momentum but got {momentum}. Please make sure the ' f'value is correct.', logger='current', level=logging.WARNING) self.momentum = momentum
[文档] def avg_func(self, averaged_param: Tensor, source_param: Tensor, steps: int) -> None: """Compute the moving average of the parameters using exponential moving average. Args: averaged_param (Tensor): The averaged parameters. source_param (Tensor): The source parameters. steps (int): The number of times the parameters have been updated. """ averaged_param.lerp_(source_param, self.momentum)
[文档]@MODELS.register_module() class MomentumAnnealingEMA(ExponentialMovingAverage): r"""Exponential moving average (EMA) with momentum annealing strategy. Args: model (nn.Module): The model to be averaged. momentum (float): The momentum used for updating ema parameter. Defaults to 0.0002. Ema's parameter are updated with the formula :math:`averaged\_param = (1-momentum) * averaged\_param + momentum * source\_param`. 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 max(momentum, gamma / (gamma + steps)) Defaults to 100. 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 = 100, interval: int = 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 linear 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 = max(self.momentum, self.gamma / (self.gamma + self.steps.item())) averaged_param.lerp_(source_param, momentum)

© Copyright 2022, mmengine contributors. Revision d480df71.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: v0.8.3
Versions
latest
stable
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.