<jupyter_start><jupyter_text>Tweaks to the OptimWrapper to handle an array of lrs/wds/... We will modify OptimWrapper so that it accepts lists of hyperparameters (learning rate, weight decay, momentum, beta, alpha). We will use this by first defining a set of groups of layers and then defining one hyperparameter value for each group.An example of this is _Differential Learning_ (lately [renamed](https://twitter.com/jeremyphoward/status/954459083788058624) to _Discriminative Fine Tuning_) which consists of using different learning rates on different layers during training.<jupyter_code>#export
def requires_grad(l, b=None):
    ps = list(l.parameters())
    if not ps: return None
    if b is None: return ps[0].requires_grad
    for p in ps: p.requires_grad=b

def trainable_params(m): return filter(lambda p: p.requires_grad, m.parameters())
#export
class OptimWrapper():
    "Basic wrapper around an optimizer to simplify HP changes"
    def __init__(self, opt:optim.Optimizer, wd:Floats=0., true_wd:bool=False):
        self.opt,self.true_wd = opt,true_wd
        self.opt_keys = list(self.opt.param_groups[0].keys())
        self.opt_keys.remove('params')
        self.read_defaults()
        self.wd = wd
        
    @classmethod
    def create(cls, opt_fn, lr, layer_groups, **kwargs):
        opt = opt_fn([{'params': trainable_params(l), 'lr':0} for l in layer_groups])
        opt = cls(opt, **kwargs)
        opt.lr = listify(lr, layer_groups)
        return opt
    
    def __repr__(self) -> str:
        return f'OptimWrapper over {repr(self.opt)}.\nTrue weight decay: {self.true_wd}'

    #Pytorch optimizer methods
    def step(self):
        # weight decay outside of optimizer step (AdamW)
        if self.true_wd:
            for lr,wd,pg in zip(self._lr,self._wd,self.opt.param_groups):
                for p in pg['params']: p.data.mul_(1 - wd*lr)
            self.set_val('weight_decay', listify(0, self._wd))
        self.opt.step()
    
    def zero_grad(self): self.opt.zero_grad()
    
    #Hyperparameters as properties
    @property
    def lr(self) -> float: return self._lr[-1]

    @lr.setter
    def lr(self, val:float): self._lr = self.set_val('lr', listify(val, self._lr))
    
    @property
    def mom(self) -> float: return self._mom[-1]

    @mom.setter
    def mom(self, val:float):
        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)
    
    @property
    def beta(self) -> float: return None if self._beta is None else self._beta[-1]

    @beta.setter
    def beta(self, val:float):
        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)
    
    @property
    def wd(self) -> float: return self._wd[-1]

    @wd.setter
    def wd(self, val:float):
        if not self.true_wd: self.set_val('weight_decay', listify(val, self._wd))
        self._wd = listify(val, self._wd)
    
    #Helper functions
    def read_defaults(self):
        "Read the values inside the optimizer for the hyper-parameters"
        self._beta = None
        if 'lr' in self.opt_keys: self._lr = self.read_val('lr')
        if 'momentum' in self.opt_keys: self._mom = self.read_val('momentum')
        if 'alpha' in self.opt_keys: self._beta = self.read_val('alpha')
        if 'betas' in self.opt_keys: self._mom,self._beta = self.read_val('betas')
        if 'weight_decay' in self.opt_keys: self._wd = self.read_val('weight_decay')
    
    def set_val(self, key:str, val):
        "Set the values inside the optimizer dictionary at the key"
        if is_tuple(val): val = [(v1,v2) for v1,v2 in zip(*val)]
        for v,pg in zip(val,self.opt.param_groups): pg[key] = v
        return val
    
    def read_val(self, key:str) -> Union[List[float],Tuple[List[float],List[float]]]:
        "Read a hyper-parameter key in the optimizer dictionary."
        val = [pg[key] for pg in self.opt.param_groups]
        if is_tuple(val[0]): val = [o[0] for o in val], [o[1] for o in val]
        return val
#export
flatten_model=lambda l: sum(map(flatten_model,l.children()),[]) if len(list(l.children())) else [l]
def first_layer(m): return flatten_model(m)[0]

def split_model_idx(model, idxs):
    "Split the model according to the indices in [idxs]"
    layers = flatten_model(model)
    if idxs[0] != 0: idxs = [0] + idxs
    if idxs[-1] != len(layers): idxs.append(len(layers))
    return [nn.Sequential(*layers[i:j]) for i,j in zip(idxs[:-1],idxs[1:])]

def split_model(model, splits, want_idxs=False):
    "Split the model according to the layers in [splits]"
    layers = flatten_model(model)
    idxs = [layers.index(first_layer(s)) for s in listify(splits)]
    res = split_model_idx(model, idxs)
    return (res,idxs) if want_idxs else res<jupyter_output><empty_output><jupyter_text>We will now try our Wrapper with the Darknet neural net architecture.<jupyter_code>model = Darknet([1, 2, 2, 2, 2], num_classes=2, nf=16)<jupyter_output><empty_output><jupyter_text>We are now going to split the model in layers so that we can try our OptimWrapper with more than one value per hyperparameter. In particular, we are going to split the 18 layers in the model into three groups: 0-5, 6-9 and 10-18. Afterwards, when setting the value for our hyperparameters, we will need to define a different value for each of these groups.<jupyter_code>splits = [model.layers[9], model.layers[15]]
layer_groups,idxs = split_model(model, splits, want_idxs=True)
lrs = np.array([1e-3,1e-2,0.1])
tst_opt = OptimWrapper(optim.SGD([{'params':l.parameters(), 'lr':lr} for l,lr in zip(layer_groups, lrs)]))
idxs<jupyter_output><empty_output><jupyter_text>We can set the hyperparameters in two ways:1. optimizer_object.\_hyperparameter = [val1, val2 ..., valn] (n = number of layer groups)2. optimizer_object.hyperparameter = valIf we chose to set it in way 1, we must specify a number of values exactly equal to the number of layer groups. If we chose to set it in way 2, the chosen value will be repeated for all layer groups.<jupyter_code>tst_opt.lr, tst_opt._lr
tst_opt.wd, tst_opt._wd
# expect exception
# tst_opt.wd = [0.1,0.1]
# expect exception
# tst_opt.lr = [0.1,0.1]<jupyter_output><empty_output><jupyter_text>Now let's tweak the learner to handle this. We have added some important functions into the Learner class. We have built *create_opt* that creates the optimizer with the wrapper, _split_ that splits our model into layer groups, *freeze_to*, _freeze_ and _unfreeze_ that allow us to freeze and unfreeze parts of the network or even the whole network. Add gradient clipping at this stage in a callback.<jupyter_code>#export
bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)

def set_bn_eval(m):
    for l in m.children():
        if isinstance(l, bn_types) and not next(l.parameters()).requires_grad:
            l.eval()
        set_bn_eval(l)

@dataclass
class BnFreeze(Callback):
    learn:Learner
    def on_epoch_begin(self, **kwargs): set_bn_eval(self.learn.model)

def even_mults(start, stop, n):
    mult = stop/start
    step = mult**(1/(n-1))
    return np.array([start*(step**i) for i in range(n)])
#export
@dataclass
class Learner():
    "Object that wraps together some data, a model, a loss function and an optimizer"
    data:DataBunch
    model:nn.Module
    opt_fn:Callable=AdamW
    loss_fn:Callable=F.cross_entropy
    metrics:Collection[Callable]=None
    true_wd:bool=True
    wd:Floats=1e-6
    train_bn:bool=True
    path:str = 'models'
    callback_fns:Collection[Callable]=None
    callbacks:Collection[Callback]=field(default_factory=list)
    layer_groups:Collection[nn.Module]=None
    def __post_init__(self):
        self.path = Path(self.path)
        self.path.mkdir(parents=True, exist_ok=True)
        self.model = self.model.to(self.data.device)
        self.metrics=listify(self.metrics)
        if not self.layer_groups: self.layer_groups = [self.model]
        self.callbacks = listify(self.callbacks)
        self.callback_fns = [Recorder] + listify(self.callback_fns)

    def lr_range(self, lr):
        if not isinstance(lr,slice): return lr
        if lr.start: res = even_mults(lr.start, lr.stop, len(self.layer_groups))
        else: res = [lr.stop/3]*(len(self.layer_groups)-1) + [lr.stop]
        return np.array(res)
        
    def fit(self, epochs:int, lr:Union[Floats,slice], wd:Floats=None, callbacks:Collection[Callback]=None):
        lr = self.lr_range(lr)
        if wd is None: wd = self.wd
        self.create_opt(lr, wd)
        callbacks = [cb(self) for cb in self.callback_fns] + listify(callbacks)
        fit(epochs, self.model, self.loss_fn, self.opt, self.data, metrics=self.metrics,
            callbacks=self.callbacks+callbacks)

    def create_opt(self, lr:Floats, wd:Floats=0.):
        self.opt = OptimWrapper.create(self.opt_fn, lr, self.layer_groups, wd=wd, true_wd=self.true_wd)
        
    def split(self, split_on):
        if isinstance(split_on,Callable): split_on = split_on(self.model)
        self.layer_groups = split_model(self.model, split_on)

    def freeze_to(self, n):
        for g in self.layer_groups[:n]:
            for l in g:
                if not self.train_bn or not isinstance(l, bn_types): requires_grad(l, False)
        for g in self.layer_groups[n:]: requires_grad(g, True)

    def freeze(self):
        assert(len(self.layer_groups)>1)
        self.freeze_to(-1)
        
    def unfreeze(self): self.freeze_to(0)
        
    def save(self, name): torch.save(self.model.state_dict(), self.path/f'{name}.pth')
    def load(self, name): self.model.load_state_dict(torch.load(self.path/f'{name}.pth'))
#export
def fit_one_cycle(learn:Learner, cyc_len:int, max_lr:float, moms:Tuple[float,float]=(0.95,0.85),
                  div_factor:float=10., pct_start:float=0.5, pct_end:float=0.3, wd:float=0.):
    "Fits a model following the 1cycle policy"
    max_lr = learn.lr_range(max_lr)
    cbs = [OneCycleScheduler(learn, max_lr, moms=moms, div_factor=div_factor,
                             pct_start=pct_start, pct_end=pct_end)]
    learn.fit(cyc_len, max_lr, wd=wd, callbacks=cbs)

Learner.fit_one_cycle = fit_one_cycle
Learner.lr_find = lr_find
train_ds = FilesDataset.from_folder(PATH/'train', classes=['airplane','dog'])
valid_ds = FilesDataset.from_folder(PATH/'test', classes=['airplane','dog'])
data = DataBunch.create(train_ds, valid_ds, bs=bs, train_tfm=train_tfms, valid_tfm=valid_tfms, num_workers=0, dl_tfms=cifar_norm)
len(data.train_dl), len(data.valid_dl)
model = Darknet([1, 2, 2, 2, 2], num_classes=2, nf=16)
learn = Learner(data, model, true_wd=True)
learn.metrics = [accuracy]
learn.split((model.layers[9],model.layers[15]))
learn.freeze()
l1,l2 = model.layers[0],model.layers[-1]; l1,l2
requires_grad(l1[0]),requires_grad(l1[1]),requires_grad(l2)<jupyter_output><empty_output><jupyter_text>As we know, setting a hyperparameter to a list of length different from the number of layer groups will return an error.<jupyter_code># expect exception
# learn.fit(1, [0.1,0.1], wd=1e-2)
learn.fit_one_cycle(1, slice(1e-1))
learn.fit(1, 1e-1, wd=1e-2)
learn.opt._lr
learn.fit(1, slice(1e-3,1e-1), wd=1e-2)
learn.opt._lr
learn.unfreeze()
requires_grad(l1[0]),requires_grad(l1[1]),requires_grad(l2)
learn.fit(1, lrs, wd=1e-2)
model = Darknet([1, 2, 2, 2, 2], num_classes=2, nf=16)
learn = Learner(data, model)
learn.metrics = [accuracy]
learn.fit(1,0.1)
learn.opt._lr
model = Darknet([1, 2, 2, 2, 2], num_classes=2, nf=16)
learn = Learner(data, model)
learn.metrics = [accuracy]
learn.layer_groups = split_model_idx(model, idxs)
learn.fit(1, lrs, wd=[1e-4,1e-3,1e-2])
learn.opt._wd
# exception expected
# learn.fit(1, lrs, wd=[1e-4,1e-3])<jupyter_output><empty_output><jupyter_text>LRs and WDs are the easiest to pass through the Learner, but if a Callback sets an array of moms or betas, the OptimWrapper will handle them as discriminative moms/betas.  See how it fits with the other callbacks<jupyter_code>model = Darknet([1, 2, 2, 2, 2], num_classes=2, nf=16)
learn = Learner(data, model)
learn.metrics = [accuracy]
learn.layer_groups = split_model_idx(model, idxs)
lr_find(learn, start_lr=lrs/1000, end_lr=lrs*100)
learn.recorder.plot()
sched = OneCycleScheduler(learn, lrs)
learn.fit(1, lrs, callbacks=[sched])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12815827109486733, 'fraction_numerical': 0.01668520578420467, 'mean_word_length': 3.85796989579313, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 36, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '370917', 'license': 'Apache-2.0', 'n_tokens_mistral': 4575, 'n_tokens_neox': 4453, 'n_words': 1124, 'path': 'dev_nb/004a_discriminative_lr.ipynb', 'repo_name': 'lesscomfortable/fastai_v1', 'size': '12586'}