Creating a new learner fit function

maxmatical · March 3, 2020, 2:36am

I’m trying to create a knowledge distillation/self training process with a teacher and student model, and I would like to create a fit_distill function based off the fit function that takes the model outputs from a teacher and student model and calculates the distillation loss, which also requires changing the loss_batch function. I have an idea on how to create the fit_distill function, but once that’s done, how can I call learn.fit_distill? Is there a way to add a function into the learner class as a method?

muellerzr · March 3, 2020, 3:34am

Look at the source code for fit_one_cycle. You’ll see that it actually takes in a learner, so you could do fit_one_cycle(learn, 5, 1e-3)

github.com

fastai/fastai/blob/master/fastai/callbacks/one_cycle.py

"Supports 1-Cycle style training"
from ..core import *
from ..callback import *
from ..basic_train import Learner,LearnerCallback

__all__ = ['OneCycleScheduler']

class OneCycleScheduler(LearnerCallback):
    "Manage 1-Cycle style training as outlined in Leslie Smith's [paper](https://arxiv.org/pdf/1803.09820.pdf)."
    def __init__(self, learn:Learner, lr_max:float, moms:Floats=(0.95,0.85), div_factor:float=25., pct_start:float=0.3,
                 final_div:float=None, tot_epochs:int=None, start_epoch:int=None):
        super().__init__(learn)
        self.lr_max,self.div_factor,self.pct_start,self.final_div = lr_max,div_factor,pct_start,final_div
        if self.final_div is None: self.final_div = div_factor*1e4
        self.moms=tuple(listify(moms,2))
        if is_listy(self.lr_max): self.lr_max = np.array(self.lr_max)
        self.start_epoch, self.tot_epochs = start_epoch, tot_epochs

    def steps(self, *steps_cfg:StartOptEnd):
        "Build anneal schedule for all of the parameters."

This file has been truncated. show original

An easier one to read and understand may be the flat_cos_anneal fit function:

github.com

fastai/fastai/blob/master/fastai/train.py#L25


def fit_one_cycle(learn:Learner, cyc_len:int, max_lr:Union[Floats,slice]=defaults.lr,
                  moms:Tuple[float,float]=(0.95,0.85), div_factor:float=25., pct_start:float=0.3, final_div:float=None,
                  wd:float=None, callbacks:Optional[CallbackList]=None, tot_epochs:int=None, start_epoch:int=None)->None:
    "Fit a model following the 1cycle policy."
    max_lr = learn.lr_range(max_lr)
    callbacks = listify(callbacks)
    callbacks.append(OneCycleScheduler(learn, max_lr, moms=moms, div_factor=div_factor, pct_start=pct_start,
                                       final_div=final_div, tot_epochs=tot_epochs, start_epoch=start_epoch))
    learn.fit(cyc_len, max_lr, wd=wd, callbacks=callbacks)


def fit_fc(learn:Learner, tot_epochs:int=1, lr:float=defaults.lr,  moms:Tuple[float,float]=(0.95,0.85), start_pct:float=0.72,
                  wd:float=None, callbacks:Optional[CallbackList]=None)->None:
    "Fit a model with Flat Cosine Annealing"
    max_lr = learn.lr_range(lr)
    callbacks = listify(callbacks)
    callbacks.append(FlatCosAnnealScheduler(learn, lr, moms=moms, start_pct=start_pct, tot_epochs=tot_epochs))
    learn.fit(tot_epochs, max_lr, wd=wd, callbacks=callbacks)


def lr_find(learn:Learner, start_lr:Floats=1e-7, end_lr:Floats=10, num_it:int=100, stop_div:bool=True, wd:float=None):
    "Explore lr from `start_lr` to `end_lr` over `num_it` iterations in `learn`. If `stop_div`, stops when loss diverges."
    start_lr = learn.lr_range(start_lr)