Shortcuts

注意

您正在阅读 MMEditing 0.x。 MMEditing 0.x 会在 2022 年末开始逐步停止维护,建议您及时升级到 MMEditing 1.0 版本,享受由 OpenMMLab 2.0 带来的更多新特性和更佳的性能表现。阅读 MMEditing 1.0 的发版日志代码文档 以了解更多。

mmedit.core.optimizer.builder 源代码

# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.runner import build_optimizer


[文档]def build_optimizers(model, cfgs): """Build multiple optimizers from configs. If `cfgs` contains several dicts for optimizers, then a dict for each constructed optimizers will be returned. If `cfgs` only contains one optimizer config, the constructed optimizer itself will be returned. For example, 1) Multiple optimizer configs: .. code-block:: python optimizer_cfg = dict( model1=dict(type='SGD', lr=lr), model2=dict(type='SGD', lr=lr)) The return dict is ``dict('model1': torch.optim.Optimizer, 'model2': torch.optim.Optimizer)`` 2) Single optimizer config: .. code-block:: python optimizer_cfg = dict(type='SGD', lr=lr) The return is ``torch.optim.Optimizer``. Args: model (:obj:`nn.Module`): The model with parameters to be optimized. cfgs (dict): The config dict of the optimizer. Returns: dict[:obj:`torch.optim.Optimizer`] | :obj:`torch.optim.Optimizer`: The initialized optimizers. """ optimizers = {} if hasattr(model, 'module'): model = model.module # determine whether 'cfgs' has several dicts for optimizers is_dict_of_dict = True for key, cfg in cfgs.items(): if not isinstance(cfg, dict): is_dict_of_dict = False if is_dict_of_dict: for key, cfg in cfgs.items(): cfg_ = cfg.copy() module = getattr(model, key) optimizers[key] = build_optimizer(module, cfg_) return optimizers return build_optimizer(model, cfgs)