fl_sim.optimizers.get_optimizer

fl_sim.optimizers.get_optimizer#

fl_sim.optimizers.get_optimizer(optimizer_name: str | type, params: Iterable[dict | Parameter], config: Any) → Optimizer[源代码]#

Get optimizer by name.

参数:
  • optimizer_name (Union[str, type]) -- Optimizer name or class

  • params (Iterable[Union[dict, torch.nn.parameter.Parameter]]) -- Parameters to be optimized

  • config (Any) -- Config for optimizer. Should be a dict or a class with attributes which can be accessed by config.attr.

返回:

Instance of the given optimizer.

返回类型:

torch.optim.Optimizer

示例

import torch

model = torch.nn.Linear(10, 1)
optimizer = get_optimizer("SGD", model.parameters(), {"lr": 1e-2})  # PyTorch built-in
optimizer = get_optimizer("yogi", model.parameters(), {"lr": 1e-2})  # from pytorch_optimizer
optimizer = get_optimizer("FedPD_SGD", model.parameters(), {"lr": 1e-2})  # federated