facebookresearch / facebookresearch/detectron2
Add support to select an optimizer
- Dominant language
- Python
- Stars
- 34.7k
- Forks
- 7.9k
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Other torchvision optimizers can be used to train Detectron2 models.
## Motivation & Examples
SGD convergence is pretty slow. When training Detectron2 models from scratch we have to set very low learning rate in order to avoid `FLoatingPointError` (#4804). Other optimizers such as Adam, RMSprop, etc. can speed up training.
It will be passed from `cfg` as `cfg.SOLVER.OPTIMIZER = "SGD"` and then parsed when building an optimizer:
```python
def build_optimizer(cfg: CfgNode, model: torch.nn.Module) -> torch.optim.Optimizer:
"""
Build an optimizer from config.
"""
params = get_default_optimizer_params(
model,
base_lr=cfg.SOLVER.BASE_LR,
weight_decay_norm=cfg.SOLVER.WEIGHT_DECAY_NORM,
bias_lr_factor=cfg.SOLVER.BIAS_LR_FACTOR,
weight_decay_bias=cfg.SOLVER.WEIGHT_DECAY_BIAS,
)
sgd_args = {
"params": params,
"lr": cfg.SOLVER.BASE_LR,
"momentum": cfg.SOLVER.MOMENTUM,
"nesterov": cfg.SOLVER.NESTEROV,
"weight_decay": cfg.SOLVER.WEIGHT_DECAY,
}
if TORCH_VERSION >= (1, 12):
sgd_args["foreach"] = True
optimizer = getattr(orch.optim, cfg.SOLVER.OPTIMIZER)
if optimizer is None:
raise ValueError()
return maybe_add_gradient_clipping(cfg, optimizer(**sgd_args))
```
and the default value will be set to `_C.SOLVER.NUM_DECAYS = "SGD"`
Contributor guide
Assessment
This issue has not been assessed yet.