Shortcuts

mmengine.optim.optimizer.builder 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import copy
import inspect
from typing import List, Union

import torch
import torch.nn as nn

from mmengine.config import Config, ConfigDict
from mmengine.device import is_npu_available, is_npu_support_full_precision
from mmengine.registry import OPTIM_WRAPPER_CONSTRUCTORS, OPTIMIZERS
from .optimizer_wrapper import OptimWrapper


def register_torch_optimizers() -> List[str]:
    """Register optimizers in ``torch.optim`` to the ``OPTIMIZERS`` registry.

    Returns:
        List[str]: A list of registered optimizers' name.
    """
    torch_optimizers = []
    for module_name in dir(torch.optim):
        if module_name.startswith('__'):
            continue
        _optim = getattr(torch.optim, module_name)
        if inspect.isclass(_optim) and issubclass(_optim,
                                                  torch.optim.Optimizer):
            OPTIMIZERS.register_module(module=_optim)
            torch_optimizers.append(module_name)
    return torch_optimizers


TORCH_OPTIMIZERS = register_torch_optimizers()


def register_torch_npu_optimizers() -> List[str]:
    """Register optimizers in ``torch npu`` to the ``OPTIMIZERS`` registry.

    Returns:
        List[str]: A list of registered optimizers' name.
    """
    if not is_npu_available():
        return []

    import torch_npu
    if not hasattr(torch_npu, 'optim'):
        return []

    torch_npu_optimizers = []
    for module_name in dir(torch_npu.optim):
        if module_name.startswith('__') or module_name in OPTIMIZERS:
            continue
        _optim = getattr(torch_npu.optim, module_name)
        if inspect.isclass(_optim) and issubclass(_optim,
                                                  torch.optim.Optimizer):
            OPTIMIZERS.register_module(module=_optim)
            torch_npu_optimizers.append(module_name)
    return torch_npu_optimizers


NPU_OPTIMIZERS = register_torch_npu_optimizers()


def register_dadaptation_optimizers() -> List[str]:
    """Register optimizers in ``dadaptation`` to the ``OPTIMIZERS`` registry.

    Returns:
        List[str]: A list of registered optimizers' name.
    """
    dadaptation_optimizers = []
    try:
        import dadaptation
    except ImportError:
        pass
    else:
        for module_name in ['DAdaptAdaGrad', 'DAdaptAdam', 'DAdaptSGD']:
            _optim = getattr(dadaptation, module_name)
            if inspect.isclass(_optim) and issubclass(_optim,
                                                      torch.optim.Optimizer):
                OPTIMIZERS.register_module(module=_optim)
                dadaptation_optimizers.append(module_name)
    return dadaptation_optimizers


DADAPTATION_OPTIMIZERS = register_dadaptation_optimizers()


def register_lion_optimizers() -> List[str]:
    """Register Lion optimizer to the ``OPTIMIZERS`` registry.

    Returns:
        List[str]: A list of registered optimizers' name.
    """
    optimizers = []
    try:
        from lion_pytorch import Lion
    except ImportError:
        pass
    else:
        OPTIMIZERS.register_module(module=Lion)
        optimizers.append('Lion')
    return optimizers


LION_OPTIMIZERS = register_lion_optimizers()


def register_sophia_optimizers() -> List[str]:
    """Register Sophia optimizer to the ``OPTIMIZERS`` registry.

    Returns:
        List[str]: A list of registered optimizers' name.
    """
    optimizers = []
    try:
        import Sophia
    except ImportError:
        pass
    else:
        for module_name in dir(Sophia):
            _optim = getattr(Sophia, module_name)
            if inspect.isclass(_optim) and issubclass(_optim,
                                                      torch.optim.Optimizer):
                OPTIMIZERS.register_module(module=_optim)
                optimizers.append(module_name)
    return optimizers


SOPHIA_OPTIMIZERS = register_sophia_optimizers()


def register_deepspeed_optimizers() -> List[str]:
    """Register optimizers in ``deepspeed`` to the ``OPTIMIZERS`` registry.

    Returns:
        List[str]: A list of registered optimizers' name.
    """
    deepspeed_optimizers = []
    try:
        import deepspeed  # noqa: F401
    except ImportError:
        pass
    else:
        from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
        from deepspeed.ops.lamb import FusedLamb
        from deepspeed.runtime.fp16.onebit import (OnebitAdam, OnebitLamb,
                                                   ZeroOneAdam)

        OPTIMIZERS.register_module(module=DeepSpeedCPUAdam)
        deepspeed_optimizers.append('DeepSpeedCPUAdam')
        OPTIMIZERS.register_module(module=FusedAdam)
        deepspeed_optimizers.append('FusedAdam')
        OPTIMIZERS.register_module(module=FusedLamb)
        deepspeed_optimizers.append('FusedLamb')
        OPTIMIZERS.register_module(module=OnebitAdam)
        deepspeed_optimizers.append('OnebitAdam')
        OPTIMIZERS.register_module(module=OnebitLamb)
        deepspeed_optimizers.append('OnebitLamb')
        OPTIMIZERS.register_module(module=ZeroOneAdam)
        deepspeed_optimizers.append('ZeroOneAdam')

    return deepspeed_optimizers


DEEPSPEED_OPTIMIZERS = register_deepspeed_optimizers()


[文档]def build_optim_wrapper(model: nn.Module, cfg: Union[dict, Config, ConfigDict]) -> OptimWrapper: """Build function of OptimWrapper. If ``constructor`` is set in the ``cfg``, this method will build an optimizer wrapper constructor, and use optimizer wrapper constructor to build the optimizer wrapper. If ``constructor`` is not set, the ``DefaultOptimWrapperConstructor`` will be used by default. Args: model (nn.Module): Model to be optimized. cfg (dict): Config of optimizer wrapper, optimizer constructor and optimizer. Returns: OptimWrapper: The built optimizer wrapper. """ optim_wrapper_cfg = copy.deepcopy(cfg) constructor_type = optim_wrapper_cfg.pop('constructor', 'DefaultOptimWrapperConstructor') paramwise_cfg = optim_wrapper_cfg.pop('paramwise_cfg', None) # Since the current generation of NPU(Ascend 910) only supports # mixed precision training, here we turn on mixed precision # to make the training normal if is_npu_available() and not is_npu_support_full_precision(): optim_wrapper_cfg['type'] = 'AmpOptimWrapper' optim_wrapper_constructor = OPTIM_WRAPPER_CONSTRUCTORS.build( dict( type=constructor_type, optim_wrapper_cfg=optim_wrapper_cfg, paramwise_cfg=paramwise_cfg)) optim_wrapper = optim_wrapper_constructor(model) return optim_wrapper

© Copyright 2022, mmengine contributors. Revision b2295a25.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: v0.8.1
Versions
latest
stable
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.