Spaces:
Build error
Build error
File size: 4,864 Bytes
c426e13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
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:
# Separate classifier parameters from all others
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
|