Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,660 Bytes
78e32cc |
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 |
###
# Author: Kai Li
# Date: 2021-06-20 00:21:33
# LastEditors: Please set LastEditors
# LastEditTime: 2022-05-27 11:19:51
###
from torch.optim.optimizer import Optimizer
from torch.optim import Adam, RMSprop, SGD, Adadelta, Adagrad, Adamax, AdamW, ASGD
from torch_optimizer import (
AccSGD,
AdaBound,
AdaMod,
DiffGrad,
Lamb,
NovoGrad,
PID,
QHAdam,
QHM,
RAdam,
SGDW,
Yogi,
Ranger,
RangerQH,
RangerVA,
)
__all__ = [
"AccSGD",
"AdaBound",
"AdaMod",
"DiffGrad",
"Lamb",
"NovoGrad",
"PID",
"QHAdam",
"QHM",
"RAdam",
"SGDW",
"Yogi",
"Ranger",
"RangerQH",
"RangerVA",
"Adam",
"RMSprop",
"SGD",
"Adadelta",
"Adagrad",
"Adamax",
"AdamW",
"ASGD",
"make_optimizer",
"get",
]
def make_optimizer(params, optim_name="adam", **kwargs):
"""
Args:
params (iterable): Output of `nn.Module.parameters()`.
optimizer (str or :class:`torch.optim.Optimizer`): Identifier understood
by :func:`~.get`.
**kwargs (dict): keyword arguments for the optimizer.
Returns:
torch.optim.Optimizer
Examples
>>> from torch import nn
>>> model = nn.Sequential(nn.Linear(10, 10))
>>> optimizer = make_optimizer(model.parameters(), optimizer='sgd',
>>> lr=1e-3)
"""
return get(optim_name)(params, **kwargs)
def register_optimizer(custom_opt):
"""Register a custom opt, gettable with `optimzers.get`.
Args:
custom_opt: Custom optimizer to register.
"""
if (
custom_opt.__name__ in globals().keys()
or custom_opt.__name__.lower() in globals().keys()
):
raise ValueError(
f"Activation {custom_opt.__name__} already exists. Choose another name."
)
globals().update({custom_opt.__name__: custom_opt})
def get(identifier):
"""Returns an optimizer function from a string. Returns its input if it
is callable (already a :class:`torch.optim.Optimizer` for example).
Args:
identifier (str or Callable): the optimizer identifier.
Returns:
:class:`torch.optim.Optimizer` or None
"""
if isinstance(identifier, Optimizer):
return identifier
elif isinstance(identifier, str):
to_get = {k.lower(): v for k, v in globals().items()}
cls = to_get.get(identifier.lower())
if cls is None:
raise ValueError(f"Could not interpret optimizer : {str(identifier)}")
return cls
raise ValueError(f"Could not interpret optimizer : {str(identifier)}")
|