mirror of
https://github.com/jantic/DeOldify.git
synced 2026-09-01 04:58:26 +08:00
Jason Antic
92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
from .imports import *
|
|
from .torch_imports import *
|
|
from .core import *
|
|
|
|
def opt_params(parm, lr, wd):
|
|
return {'params': chain_params(parm), 'lr':lr, 'weight_decay':wd}
|
|
|
|
class LayerOptimizer():
|
|
def __init__(self, opt_fn, layer_groups, lrs, wds=None):
|
|
if not isinstance(layer_groups, (list,tuple)): layer_groups=[layer_groups]
|
|
if not isinstance(lrs, Iterable): lrs=[lrs]
|
|
if len(lrs)==1: lrs=lrs*len(layer_groups)
|
|
if wds is None: wds=0.
|
|
if not isinstance(wds, Iterable): wds=[wds]
|
|
if len(wds)==1: wds=wds*len(layer_groups)
|
|
self.layer_groups,self.lrs,self.wds = layer_groups,lrs,wds
|
|
self.opt = opt_fn(self.opt_params())
|
|
|
|
def opt_params(self):
|
|
assert(len(self.layer_groups) == len(self.lrs))
|
|
assert(len(self.layer_groups) == len(self.wds))
|
|
params = list(zip(self.layer_groups,self.lrs,self.wds))
|
|
return [opt_params(*p) for p in params]
|
|
|
|
@property
|
|
def lr(self): return self.lrs[-1]
|
|
|
|
@property
|
|
def mom(self):
|
|
if 'betas' in self.opt.param_groups[0]:
|
|
return self.opt.param_groups[0]['betas'][0]
|
|
else:
|
|
return self.opt.param_groups[0]['momentum']
|
|
|
|
def set_lrs(self, lrs):
|
|
if not isinstance(lrs, Iterable): lrs=[lrs]
|
|
if len(lrs)==1: lrs=lrs*len(self.layer_groups)
|
|
set_lrs(self.opt, lrs)
|
|
self.lrs=lrs
|
|
|
|
def set_wds_out(self, wds):
|
|
if not isinstance(wds, Iterable): wds=[wds]
|
|
if len(wds)==1: wds=wds*len(self.layer_groups)
|
|
set_wds_out(self.opt, wds)
|
|
set_wds(self.opt, [0] * len(self.layer_groups))
|
|
self.wds=wds
|
|
|
|
def set_wds(self, wds):
|
|
if not isinstance(wds, Iterable): wds=[wds]
|
|
if len(wds)==1: wds=wds*len(self.layer_groups)
|
|
set_wds(self.opt, wds)
|
|
set_wds_out(self.opt, [0] * len(self.layer_groups))
|
|
self.wds=wds
|
|
|
|
def set_mom(self,momentum):
|
|
if 'betas' in self.opt.param_groups[0]:
|
|
for pg in self.opt.param_groups: pg['betas'] = (momentum, pg['betas'][1])
|
|
else:
|
|
for pg in self.opt.param_groups: pg['momentum'] = momentum
|
|
|
|
def set_beta(self,beta):
|
|
if 'betas' in self.opt.param_groups[0]:
|
|
for pg in self.opt.param_groups: pg['betas'] = (pg['betas'][0],beta)
|
|
elif 'alpha' in self.opt.param_groups[0]:
|
|
for pg in self.opt.param_groups: pg['alpha'] = beta
|
|
|
|
def set_opt_fn(self, opt_fn):
|
|
if type(self.opt) != type(opt_fn(self.opt_params())):
|
|
self.opt = opt_fn(self.opt_params())
|
|
|
|
def zip_strict_(l, r):
|
|
assert(len(l) == len(r))
|
|
return zip(l, r)
|
|
|
|
def set_lrs(opt, lrs):
|
|
if not isinstance(lrs, Iterable): lrs=[lrs]
|
|
if len(lrs)==1: lrs=lrs*len(opt.param_groups)
|
|
for pg,lr in zip_strict_(opt.param_groups,lrs): pg['lr'] = lr
|
|
|
|
def set_wds_out(opt, wds):
|
|
if not isinstance(wds, Iterable): wds=[wds]
|
|
if len(wds)==1: wds=wds*len(opt.param_groups)
|
|
assert(len(opt.param_groups) == len(wds))
|
|
for pg,wd in zip_strict_(opt.param_groups,wds): pg['wd'] = wd
|
|
|
|
def set_wds(opt, wds):
|
|
if not isinstance(wds, Iterable): wds=[wds]
|
|
if len(wds)==1: wds=wds*len(opt.param_groups)
|
|
assert(len(opt.param_groups) == len(wds))
|
|
for pg,wd in zip_strict_(opt.param_groups,wds): pg['weight_decay'] = wd
|
|
|