Spaces:
Sleeping
Sleeping
| """ Optimizer Factory w/ Custom Weight Decay | |
| Hacked together by / Copyright 2020 Ross Wightman | |
| """ | |
| from typing import Optional | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from .timm.adafactor import Adafactor | |
| from .timm.adahessian import Adahessian | |
| from .timm.adamp import AdamP | |
| from .timm.lookahead import Lookahead | |
| from .timm.nadam import Nadam | |
| from .timm.novograd import NovoGrad | |
| from .timm.nvnovograd import NvNovoGrad | |
| from .timm.radam import RAdam | |
| from .timm.rmsprop_tf import RMSpropTF | |
| from .timm.sgdp import SGDP | |
| from .timm.adabelief import AdaBelief | |
| try: | |
| from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD | |
| has_apex = True | |
| except ImportError: | |
| has_apex = False | |
| def add_weight_decay(model, weight_decay=1e-5, skip_list=()): | |
| decay = [] | |
| no_decay = [] | |
| for name, param in model.named_parameters(): | |
| if not param.requires_grad: | |
| continue # frozen weights | |
| if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list: | |
| no_decay.append(param) | |
| else: | |
| decay.append(param) | |
| return [ | |
| {"params": no_decay, "weight_decay": 0.}, | |
| {"params": decay, "weight_decay": weight_decay}] | |
| def optimizer_kwargs(args, lr_weight): | |
| """ args/argparse to kwargs helper | |
| Convert optimizer args in argparse args or args like object to keyword args for updated create fn. | |
| """ | |
| kwargs = dict( | |
| optimizer_name=args.opt, | |
| learning_rate=args.lr_base*args.batch_size/128*lr_weight, | |
| weight_decay=args.weight_decay, | |
| momentum=args.momentum) | |
| if getattr(args, "opt_eps", None) is not None: | |
| kwargs["eps"] = args.opt_eps | |
| if getattr(args, "opt_betas", None) is not None: | |
| kwargs["betas"] = args.opt_betas | |
| if getattr(args, "opt_args", None) is not None: | |
| kwargs.update(args.opt_args) | |
| return kwargs | |
| def create_optimizer(args, model, filter_bias_and_bn=True, lr_weight=1): | |
| """ Legacy optimizer factory for backwards compatibility. | |
| NOTE: Use create_optimizer_v2 for new code. | |
| """ | |
| return create_optimizer_v2( | |
| model, | |
| **optimizer_kwargs(args, lr_weight), | |
| filter_bias_and_bn=filter_bias_and_bn, | |
| ) | |
| def create_optimizer_v2( | |
| model: nn.Module, | |
| optimizer_name: str = "sgd", | |
| learning_rate: Optional[float] = None, | |
| weight_decay: float = 0., | |
| momentum: float = 0.9, | |
| filter_bias_and_bn: bool = True, | |
| **kwargs): | |
| """ Create an optimizer. | |
| TODO currently the model is passed in and all parameters are selected for optimization. | |
| For more general use an interface that allows selection of parameters to optimize and lr groups, one of: | |
| * a filter fn interface that further breaks params into groups in a weight_decay compatible fashion | |
| * expose the parameters interface and leave it up to caller | |
| Args: | |
| model (nn.Module): model containing parameters to optimize | |
| optimizer_name: name of optimizer to create | |
| learning_rate: initial learning rate | |
| weight_decay: weight decay to apply in optimizer | |
| momentum: momentum for momentum based optimizers (others may use betas via kwargs) | |
| filter_bias_and_bn: filter out bias, bn and other 1d params from weight decay | |
| **kwargs: extra optimizer specific kwargs to pass through | |
| Returns: | |
| Optimizer | |
| """ | |
| opt_lower = optimizer_name.lower() | |
| if weight_decay and filter_bias_and_bn: | |
| skip = {} | |
| if hasattr(model, "no_weight_decay"): | |
| skip = model.no_weight_decay() | |
| parameters = add_weight_decay(model, weight_decay, skip) | |
| weight_decay = 0. | |
| else: | |
| parameters = model.parameters() | |
| if "fused" in opt_lower: | |
| assert has_apex and torch.cuda.is_available(), "APEX and CUDA required for fused optimizers" | |
| opt_args = dict(lr=learning_rate, weight_decay=weight_decay, **kwargs) | |
| opt_split = opt_lower.split("_") | |
| opt_lower = opt_split[-1] | |
| if opt_lower == "sgd" or opt_lower == "nesterov": | |
| opt_args.pop("eps", None) | |
| optimizer = optim.SGD(parameters, momentum=momentum, nesterov=True, **opt_args) | |
| elif opt_lower == "momentum": | |
| opt_args.pop("eps", None) | |
| optimizer = optim.SGD(parameters, momentum=momentum, nesterov=False, **opt_args) | |
| elif opt_lower == "adam": | |
| optimizer = optim.Adam(parameters, **opt_args) | |
| elif opt_lower == "adabelief": | |
| optimizer = AdaBelief(parameters, rectify=False, **opt_args) | |
| elif opt_lower == "adamw": | |
| optimizer = optim.AdamW(parameters, lr=learning_rate, weight_decay=weight_decay) | |
| elif opt_lower == "nadam": | |
| optimizer = Nadam(parameters, **opt_args) | |
| elif opt_lower == "radam": | |
| optimizer = RAdam(parameters, **opt_args) | |
| elif opt_lower == "adamp": | |
| optimizer = AdamP(parameters, wd_ratio=0.01, nesterov=True, **opt_args) | |
| elif opt_lower == "sgdp": | |
| optimizer = SGDP(parameters, momentum=momentum, nesterov=True, **opt_args) | |
| elif opt_lower == "adadelta": | |
| optimizer = optim.Adadelta(parameters, **opt_args) | |
| elif opt_lower == "adafactor": | |
| if not learning_rate: | |
| opt_args["lr"] = None | |
| optimizer = Adafactor(parameters, **opt_args) | |
| elif opt_lower == "adahessian": | |
| optimizer = Adahessian(parameters, **opt_args) | |
| elif opt_lower == "rmsprop": | |
| optimizer = optim.RMSprop(parameters, alpha=0.9, momentum=momentum, **opt_args) | |
| elif opt_lower == "rmsproptf": | |
| optimizer = RMSpropTF(parameters, alpha=0.9, momentum=momentum, **opt_args) | |
| elif opt_lower == "novograd": | |
| optimizer = NovoGrad(parameters, **opt_args) | |
| elif opt_lower == "nvnovograd": | |
| optimizer = NvNovoGrad(parameters, **opt_args) | |
| elif opt_lower == "fusedsgd": | |
| opt_args.pop("eps", None) | |
| optimizer = FusedSGD(parameters, momentum=momentum, nesterov=True, **opt_args) | |
| elif opt_lower == "fusedmomentum": | |
| opt_args.pop("eps", None) | |
| optimizer = FusedSGD(parameters, momentum=momentum, nesterov=False, **opt_args) | |
| elif opt_lower == "fusedadam": | |
| optimizer = FusedAdam(parameters, adam_w_mode=False, **opt_args) | |
| elif opt_lower == "fusedadamw": | |
| optimizer = FusedAdam(parameters, adam_w_mode=True, **opt_args) | |
| elif opt_lower == "fusedlamb": | |
| optimizer = FusedLAMB(parameters, **opt_args) | |
| elif opt_lower == "fusednovograd": | |
| opt_args.setdefault("betas", (0.95, 0.98)) | |
| optimizer = FusedNovoGrad(parameters, **opt_args) | |
| else: | |
| assert False and "Invalid optimizer" | |
| raise ValueError | |
| if len(opt_split) > 1: | |
| if opt_split[0] == "lookahead": | |
| optimizer = Lookahead(optimizer) | |
| return optimizer | |