def wd(self, val: float) -> None:
"Set weight decay."
if not self.true_wd:
self.set_val('weight_decay',
listify(val, self._wd),
bn_groups=self.bn_wd)
self._wd = listify(val, self._wd)
def beta(self, val: float) -> None:
"Set beta (or alpha as makes sense for given optimizer)."
if val is None: return
if 'betas' in self.opt_keys:
self.set_val('betas', (self._mom, listify(val, self._beta)))
elif 'alpha' in self.opt_keys:
self.set_val('alpha', listify(val, self._beta))
self._beta = listify(val, self._beta)
def set_stat(self, name: str, value: Union[float,
Collection[float]]) -> None:
if name in ['lr', 'mom', 'beta', 'wd']:
setattr(self, name, value)
else:
val = listify(value, self._stats[name])
self.set_val(name, val)
self._stats[name] = val
def create(cls,
opt_func: object,
lr: object,
layer_groups: object,
wd: object = 0.,
true_wd: object = False,
bn_wd: object = True) -> object:
"Create an `optim.Optimizer` from `opt_func` with `lr`. Set lr on `layer_groups`."
opt: object = opt_func
opt = cls(opt, wd=wd, true_wd=true_wd, bn_wd=bn_wd)
opt.lr, opt.opt_func = listify(lr, layer_groups), opt_func
return opt
def step(self, closure=None) -> None:
"Set weight decay and step optimizer."
# weight decay outside of optimizer step (AdamW)
if self.true_wd:
lr: List[Any]
for lr, wd, pg1, pg2 in zip(self._lr, self._wd,
self.opt.param_groups[::2],
self.opt.param_groups[1::2]):
for p in pg1['params']:
p.data.mul_(1 - wd * lr)
if self.bn_wd:
for p in pg2['params']:
p.data.mul_(1 - wd * lr)
self.set_val('weight_decay', listify(0, self._wd))
self.opt.step(closure)
def mom(self, val: float) -> None:
if 'momentum' in self.opt_keys:
self.set_val('momentum', listify(val, self._mom))
elif 'betas' in self.opt_keys:
self.set_val('betas', (listify(val, self._mom), self._beta))
self._mom = listify(val, self._mom)
def lr(self, val: float) -> None:
self._lr = self.set_val('lr', listify(val, self._lr))