Spaces:
Running
on
Zero
Running
on
Zero
# Copyright (c) Facebook, Inc. and its affiliates. | |
import torch | |
from detectron2.config import CfgNode | |
from detectron2.solver import LRScheduler | |
from detectron2.solver import build_lr_scheduler as build_d2_lr_scheduler | |
from .lr_scheduler import WarmupPolyLR | |
def build_lr_scheduler(cfg: CfgNode, optimizer: torch.optim.Optimizer) -> LRScheduler: | |
""" | |
Build a LR scheduler from config. | |
""" | |
name = cfg.SOLVER.LR_SCHEDULER_NAME | |
if name == "WarmupPolyLR": | |
return WarmupPolyLR( | |
optimizer, | |
cfg.SOLVER.MAX_ITER, | |
warmup_factor=cfg.SOLVER.WARMUP_FACTOR, | |
warmup_iters=cfg.SOLVER.WARMUP_ITERS, | |
warmup_method=cfg.SOLVER.WARMUP_METHOD, | |
power=cfg.SOLVER.POLY_LR_POWER, | |
constant_ending=cfg.SOLVER.POLY_LR_CONSTANT_ENDING, | |
) | |
else: | |
return build_d2_lr_scheduler(cfg, optimizer) | |