| import cv2 |
| from apex.optimizers import FusedAdam, FusedSGD |
| from timm.optim import AdamW |
| from torch import optim |
| from torch.optim import lr_scheduler |
| from torch.optim.rmsprop import RMSprop |
| from torch.optim.adamw import AdamW |
| from torch.optim.lr_scheduler import MultiStepLR, CyclicLR |
|
|
| from training.tools.schedulers import ExponentialLRScheduler, PolyLR, LRStepScheduler |
|
|
| cv2.ocl.setUseOpenCL(False) |
| cv2.setNumThreads(0) |
|
|
|
|
| class AverageMeter(object): |
| """Computes and stores the average and current value""" |
|
|
| def __init__(self): |
| self.reset() |
|
|
| def reset(self): |
| self.val = 0 |
| self.avg = 0 |
| self.sum = 0 |
| self.count = 0 |
|
|
| def update(self, val, n=1): |
| self.val = val |
| self.sum += val * n |
| self.count += n |
| self.avg = self.sum / self.count |
|
|
| def create_optimizer(optimizer_config, model, master_params=None): |
| """Creates optimizer and schedule from configuration |
| |
| Parameters |
| ---------- |
| optimizer_config : dict |
| Dictionary containing the configuration options for the optimizer. |
| model : Model |
| The network model. |
| |
| Returns |
| ------- |
| optimizer : Optimizer |
| The optimizer. |
| scheduler : LRScheduler |
| The learning rate scheduler. |
| """ |
| if optimizer_config.get("classifier_lr", -1) != -1: |
| |
| net_params = [] |
| classifier_params = [] |
| for k, v in model.named_parameters(): |
| if not v.requires_grad: |
| continue |
| if k.find("encoder") != -1: |
| net_params.append(v) |
| else: |
| classifier_params.append(v) |
| params = [ |
| {"params": net_params}, |
| {"params": classifier_params, "lr": optimizer_config["classifier_lr"]}, |
| ] |
| else: |
| if master_params: |
| params = master_params |
| else: |
| params = model.parameters() |
|
|
| if optimizer_config["type"] == "SGD": |
| optimizer = optim.SGD(params, |
| lr=optimizer_config["learning_rate"], |
| momentum=optimizer_config["momentum"], |
| weight_decay=optimizer_config["weight_decay"], |
| nesterov=optimizer_config["nesterov"]) |
| elif optimizer_config["type"] == "FusedSGD": |
| optimizer = FusedSGD(params, |
| lr=optimizer_config["learning_rate"], |
| momentum=optimizer_config["momentum"], |
| weight_decay=optimizer_config["weight_decay"], |
| nesterov=optimizer_config["nesterov"]) |
| elif optimizer_config["type"] == "Adam": |
| optimizer = optim.Adam(params, |
| lr=optimizer_config["learning_rate"], |
| weight_decay=optimizer_config["weight_decay"]) |
| elif optimizer_config["type"] == "FusedAdam": |
| optimizer = FusedAdam(params, |
| lr=optimizer_config["learning_rate"], |
| weight_decay=optimizer_config["weight_decay"]) |
| elif optimizer_config["type"] == "AdamW": |
| optimizer = AdamW(params, |
| lr=optimizer_config["learning_rate"], |
| weight_decay=optimizer_config["weight_decay"]) |
| elif optimizer_config["type"] == "RmsProp": |
| optimizer = RMSprop(params, |
| lr=optimizer_config["learning_rate"], |
| weight_decay=optimizer_config["weight_decay"]) |
| else: |
| raise KeyError("unrecognized optimizer {}".format(optimizer_config["type"])) |
|
|
| if optimizer_config["schedule"]["type"] == "step": |
| scheduler = LRStepScheduler(optimizer, **optimizer_config["schedule"]["params"]) |
| elif optimizer_config["schedule"]["type"] == "clr": |
| scheduler = CyclicLR(optimizer, **optimizer_config["schedule"]["params"]) |
| elif optimizer_config["schedule"]["type"] == "multistep": |
| scheduler = MultiStepLR(optimizer, **optimizer_config["schedule"]["params"]) |
| elif optimizer_config["schedule"]["type"] == "exponential": |
| scheduler = ExponentialLRScheduler(optimizer, **optimizer_config["schedule"]["params"]) |
| elif optimizer_config["schedule"]["type"] == "poly": |
| scheduler = PolyLR(optimizer, **optimizer_config["schedule"]["params"]) |
| elif optimizer_config["schedule"]["type"] == "constant": |
| scheduler = lr_scheduler.LambdaLR(optimizer, lambda epoch: 1.0) |
| elif optimizer_config["schedule"]["type"] == "linear": |
| def linear_lr(it): |
| return it * optimizer_config["schedule"]["params"]["alpha"] + optimizer_config["schedule"]["params"]["beta"] |
|
|
| scheduler = lr_scheduler.LambdaLR(optimizer, linear_lr) |
|
|
| return optimizer, scheduler |
|
|
