DeOldify/fastai/callbacks/fp16.py

"Callback support for half precision (fp16) training. Increases training speed."
from ..torch_core import *
from ..callback import *
from ..basic_train import *
from torch._utils import _unflatten_dense_tensors
from torch.nn.utils import parameters_to_vector

__all__ = ['MixedPrecision']

def get_master(layer_groups:ModuleList, flat_master:bool=False) -> Tuple[List[List[Tensor]], List[List[Tensor]]]:
    "Return two lists, one for the model parameters in FP16 and one for the master parameters in FP32."
    split_params = split_no_wd_params(layer_groups)
    model_params = [[param for param in pg if param.requires_grad] for pg in split_params]
    if flat_master:
        master_params = []
        for lg in model_params:
            if len(lg) !=0 :
                mp = parameters_to_vector([param.data.float() for param in lg])
                mp = torch.nn.Parameter(mp, requires_grad=True)
                if mp.grad is None: mp.grad = mp.new(*mp.size())
                master_params.append([mp])
            else: master_params.append([])
        return model_params, master_params
    else:
        master_params = [[param.clone().float().detach() for param in lg] for lg in model_params]
        for mp in master_params:
            for param in mp: param.requires_grad = True
        return model_params, master_params

def model_g2master_g(model_params:Sequence[Tensor], master_params:Sequence[Tensor], flat_master:bool=False)->None:
    "Copy the `model_params` gradients to `master_params` for the optimizer step."
    if flat_master:
        for model_group,master_group in zip(model_params,master_params):
            if len(master_group) != 0:
                if master_group[0].grad is None: master_group[0].grad = master_group[0].data.new(*master_group[0].data.size())
                master_group[0].grad.data.copy_(parameters_to_vector([p.grad.data.float() for p in model_group]))
    else:
        for model_group,master_group in zip(model_params,master_params):
            for model, master in zip(model_group, master_group):
                if model.grad is not None:
                    if master.grad is None: master.grad = master.data.new(*master.data.size())
                    master.grad.data.copy_(model.grad.data)
                else: master.grad = None

def master2model(model_params:Sequence[Tensor], master_params:Sequence[Tensor], flat_master:bool=False)->None:
    "Copy `master_params` to `model_params`."
    if flat_master:
        for model_group,master_group in zip(model_params,master_params):
            if len(model_group) != 0:
                for model, master in zip(model_group, _unflatten_dense_tensors(master_group[0].data, model_group)):
                    model.data.copy_(master)
    else:
        for model_group,master_group in zip(model_params,master_params):
            for model, master in zip(model_group, master_group): model.data.copy_(master.data)

def grad_overflow(param_group):
    for group in param_group:
        for p in group:
            if p.grad is not None:
                s = float(p.grad.data.float().sum())
                if s == float('inf') or s == float('-inf') or s != s: return True
    return False

class MixedPrecision(LearnerCallback):
    _order = 999 #Need to run after things that could call on_backward_begin and change the loss
    "Callback that handles mixed-precision training."
    def __init__(self, learn:Learner, loss_scale:float=None, max_noskip:int=1000, dynamic:bool=True, clip:float=None,
                 flat_master:bool=False, max_scale:float=2**24):
        super().__init__(learn)
        self.flat_master,self.dynamic,self.max_noskip,self.clip,self.max_scale = flat_master,dynamic,max_noskip,clip,max_scale
        self.loss_scale = ifnone(loss_scale, 2**16 if dynamic else 512)
        self.not_min += ['model_params', 'master_params']
        assert torch.backends.cudnn.enabled, "Mixed precision training requires cudnn."
        self.opt = None

    def on_train_begin(self, **kwargs:Any)->None:
        "Prepare the master model."
        #Get a copy of the model params in FP32
        self.model_params, self.master_params = get_master(self.learn.layer_groups, self.flat_master)
        #Changes the optimizer so that the optimization step is done in FP32.
        new_opt = self.learn.opt.new_with_params(self.master_params)
        if self.opt is not None:
            self.opt.lr,self.opt.wd = self.learn.opt.lr,self.learn.opt.wd
            new_opt.load_state_dict(self.opt)
        self.learn.opt.opt = new_opt.opt
        self.noskip = 0

    def on_loss_begin(self, last_output:Tensor, **kwargs:Any) -> Tensor:
        "Convert half precision output to FP32 to avoid reduction overflow."
        return {'last_output': to_float(last_output)}

    def on_backward_begin(self, last_loss:Rank0Tensor, **kwargs:Any) -> Rank0Tensor:
        "Scale gradients up by `self.loss_scale` to prevent underflow."
        #To avoid gradient underflow, we scale the gradients
        ret_loss = last_loss * self.loss_scale
        return {'last_loss': ret_loss}

    def on_backward_end(self, **kwargs:Any)->None:
        "Convert the gradients back to FP32 and divide them by the scale."
        if self.dynamic and grad_overflow(self.model_params) and self.loss_scale > 1:
            self.loss_scale /= 2
            self.noskip = 0
            #The step will be skipped since we don't update the master grads so they are all None or zero
        else:
            model_g2master_g(self.model_params, self.master_params, self.flat_master)
            for group in self.master_params:
                for param in group:
                    if param.grad is not None: param.grad.div_(self.loss_scale)
            if self.clip is not None:
                for group in self.master_params: nn.utils.clip_grad_norm_(group, self.clip)
            if not self.dynamic: return
            self.noskip += 1
            if self.noskip >= self.max_noskip and self.loss_scale < self.max_scale:
                self.loss_scale *= 2
                self.noskip = 0

    def on_step_end(self, **kwargs:Any)->None:
        "Update the params from master to model and zero grad."
        #Zeros the gradients of the model since the optimizer is disconnected.
        self.learn.model.zero_grad()
        #Update the params from master to model.
        master2model(self.model_params, self.master_params, self.flat_master)