Shortcuts

mmengine.optim.optimizer.optimizer_wrapper_dict 源代码

# Copyright (c) OpenMMLab. All rights reserved.
from contextlib import contextmanager
from typing import Dict, Iterator, List, Optional, Tuple

import torch
import torch.nn as nn

from .optimizer_wrapper import OptimWrapper


[文档]class OptimWrapperDict(OptimWrapper): """A dictionary container of :obj:`OptimWrapper`. If runner is training with multiple optimizers, all optimizer wrappers should be managed by :obj:`OptimWrapperDict` which is built by ``CustomOptimWrapperConstructor``. ``OptimWrapperDict`` will load and save the state dictionary of all optimizer wrappers. Consider the semantic ambiguity of calling :meth:``update_params``, :meth:`backward` of all optimizer wrappers, ``OptimWrapperDict`` will not implement these methods. Examples: >>> import torch.nn as nn >>> from torch.optim import SGD >>> from mmengine.optim import OptimWrapperDict, OptimWrapper >>> model1 = nn.Linear(1, 1) >>> model2 = nn.Linear(1, 1) >>> optim_wrapper1 = OptimWrapper(SGD(model1.parameters(), lr=0.1)) >>> optim_wrapper2 = OptimWrapper(SGD(model2.parameters(), lr=0.1)) >>> optim_wrapper_dict = OptimWrapperDict(model1=optim_wrapper1, >>> model2=optim_wrapper2) Note: The optimizer wrapper contained in ``OptimWrapperDict`` can be accessed in the same way as `dict`. Args: **optim_wrappers: A dictionary of ``OptimWrapper`` instance. """ def __init__(self, **optim_wrapper_dict: OptimWrapper): for key, value in optim_wrapper_dict.items(): assert isinstance(value, OptimWrapper), ( '`OptimWrapperDict` only accept OptimWrapper instance, ' f'but got {key}: {type(value)}') self.optim_wrappers = optim_wrapper_dict
[文档] def update_params( # type: ignore self, loss: torch.Tensor, step_kwargs: Optional[Dict] = None, zero_kwargs: Optional[Dict] = None) -> None: """Update all optimizer wrappers would lead to a duplicate backward errors, and OptimWrapperDict does not know which optimizer wrapper should be updated. Therefore, this method is not implemented. The optimizer wrapper of OptimWrapperDict should be accessed and call its `update_params`. """ raise NotImplementedError('`update_params` should be called by each ' 'optimizer separately`')
[文档] def backward(self, loss: torch.Tensor, **kwargs) -> None: """Since OptimWrapperDict doesn't know which optimizer wrapper's backward method should be called (``loss_scaler`` maybe different in different :obj:AmpOptimWrapper), this method is not implemented. The optimizer wrapper of OptimWrapperDict should be accessed and call its `backward`. """ raise NotImplementedError('`backward` should be called by each ' 'optimizer separately`')
[文档] def step(self, **kwargs) -> None: """Since the backward method is not implemented, the step should not be implemented either.""" raise NotImplementedError('`step` should be called by each ' 'optimizer separately`')
[文档] def zero_grad(self, **kwargs) -> None: """Set the gradients of all optimizer wrappers to zero.""" for optim_wrapper in self.optim_wrappers.values(): optim_wrapper.zero_grad()
[文档] @contextmanager def optim_context(self, model: nn.Module): """``optim_context`` should be called by each optimizer separately.""" raise NotImplementedError( '`optim_context` should be called by each optimizer separately')
[文档] def initialize_count_status(self, model: nn.Module, cur_iter, max_iters) -> None: """Do nothing but provide unified interface for :obj:`OptimWrapper` Since ``OptimWrapperDict`` does not know the correspondence between model and optimizer wrapper. ``initialize_iter_status`` will do nothing and each optimizer wrapper should call ``initialize_iter_status`` separately. """ return
@property def param_groups(self): """Returns the parameter groups of each OptimWrapper.""" param_groups = dict() for key, value in self.optim_wrappers.items(): param_groups[key] = value.param_groups return param_groups
[文档] def get_lr(self) -> Dict[str, List[float]]: """Get the learning rate of all optimizers. Returns: Dict[str, List[float]]: Learning rate of all optimizers. """ lr_dict = dict() for name, optim_wrapper in self.optim_wrappers.items(): inner_lr_dict = optim_wrapper.get_lr() if 'base_lr' in inner_lr_dict: lr_dict[f'{name}.base_lr'] = inner_lr_dict['base_lr'] lr_dict[f'{name}.lr'] = inner_lr_dict['lr'] return lr_dict
[文档] def get_momentum(self) -> Dict[str, List[float]]: """Get the momentum of all optimizers. Returns: Dict[str, List[float]]: momentum of all optimizers. """ momentum_dict = dict() for name, optim_wrapper in self.optim_wrappers.items(): momentum_dict[f'{name}.momentum'] = optim_wrapper.get_momentum( )['momentum'] return momentum_dict
[文档] def state_dict(self) -> dict: """Get the state dictionary of all optimizer wrappers. Returns: dict: Each key-value pair in the dictionary represents the name and state dictionary of corresponding :obj:`OptimWrapper`. """ state_dict = dict() for name, optim_wrapper in self.optim_wrappers.items(): state_dict[name] = optim_wrapper.state_dict() return state_dict
[文档] def load_state_dict(self, state_dict: dict) -> None: """Load the state dictionary from the ``state_dict``. Args: state_dict (dict): Each key-value pair in `state_dict` represents the name and the state dictionary of corresponding :obj:`OptimWrapper`. """ for name, _state_dict in state_dict.items(): assert name in self.optim_wrappers, ( f'Mismatched `state_dict`! cannot found {name} in ' 'OptimWrapperDict') self.optim_wrappers[name].load_state_dict(_state_dict)
[文档] def items(self) -> Iterator[Tuple[str, OptimWrapper]]: """A generator to get the name and corresponding :obj:`OptimWrapper`""" yield from self.optim_wrappers.items()
[文档] def values(self) -> Iterator[OptimWrapper]: """A generator to get :obj:`OptimWrapper`""" yield from self.optim_wrappers.values()
[文档] def keys(self) -> Iterator[str]: """A generator to get the name of :obj:`OptimWrapper`""" yield from self.optim_wrappers.keys()
def __getitem__(self, key: str) -> OptimWrapper: assert key in self.optim_wrappers, ( f'Cannot find {key} in OptimWrapperDict, please check ' 'your optimizer constructor.') return self.optim_wrappers[key] def __contains__(self, key: str) -> bool: return key in self.optim_wrappers def __len__(self) -> int: return len(self.optim_wrappers) def __repr__(self) -> str: desc = '' for name, optim_wrapper in self.optim_wrappers.items(): desc += f'name: {name}\n' desc += repr(optim_wrapper) return desc

© Copyright 2022, mmengine contributors. Revision 85c83ba6.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
v0.10.4
v0.10.3
v0.10.2
v0.10.1
v0.10.0
v0.9.1
v0.9.0
v0.8.5
v0.8.4
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
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.