How to completely freeze the network?

ducha-aiki · January 11, 2020, 6:13pm

I feel superstupid and I must be missing something super simple.
However, here is my code (all of it)

from fastai.vision import *
mnist = untar_data(URLs.MNIST_TINY)
tfms = get_transforms(do_flip=False)
data = (ImageList.from_folder(mnist)
    .split_by_folder()          
    .label_from_folder()
    .transform(tfms, size=32)
    .databunch()
    .normalize(imagenet_stats))
learn =  cnn_learner(data, models.resnet18, metrics=accuracy, 
                 callback_fns=BnFreeze)
print (len(learn.layer_groups))

So we have 3 layer groups, I want to freeze everything

learn.freeze_to(4)
out = learn.validate()
print (out)

[1.2018166, tensor(0.0916)]

However, something is still changing:

learn.fit(1, lr=0)
out = learn.validate()
print (out)

[0.70418966, tensor(0.5408)]

ducha-aiki · January 11, 2020, 8:35pm

Ok, find the solution. BnFreeze only freezes the moving average part of the BatchNorm.
To freeze it completely, one need to also do

learn.train_bn = False

github.com

fastai/fastai/blob/master/fastai/basic_train.py#L217


    "Split the model at `split_on`."
    if isinstance(split_on,Callable): split_on = split_on(self.model)
    self.layer_groups = split_model(self.model, split_on)
    return self


def freeze_to(self, n:int)->None:
    "Freeze layers up to layer group `n`."
    if hasattr(self.model, 'reset'): self.model.reset()
    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)
    self.create_opt(defaults.lr)


def freeze(self)->None:
    "Freeze up to last layer group."
    assert(len(self.layer_groups)>1)
    self.freeze_to(-1)


def unfreeze(self):
    "Unfreeze entire model."

github.com

fastai/fastai/blob/82e8d22a6a6315f982a750789086c863a2ee5c89/fastai/train.py#L87


        "If we have `last_metrics` plot them in our pbar graph"
        if last_metrics is not None and last_metrics[0] is not None:
            rec = self.learn.recorder
            iters = range_of(rec.losses)
            val_iter = np.array(rec.nb_batches).cumsum()
            x_bounds = (0, (n_epochs - len(rec.nb_batches)) * rec.nb_batches[-1] + len(rec.losses))
            y_bounds = (0, max((max(Tensor(rec.losses)), max(Tensor(rec.val_losses)))))
            rec.pbar.update_graph([(iters, rec.losses), (val_iter, rec.val_losses)], x_bounds, y_bounds)
        return {}


class BnFreeze(LearnerCallback):
    "Freeze moving average statistics in all non-trainable batchnorm layers."
    def on_epoch_begin(self, **kwargs:Any)->None:
        "Put bn layers in eval mode just after `model.train()`."
        set_bn_eval(self.learn.model)


class GradientClipping(LearnerCallback):
    "Gradient clipping during training."
    def __init__(self, learn:Learner, clip:float = 0.):
        super().__init__(learn)
        self.clip = clip

github.com

fastai/fastai/blob/5986f1e5887038fb85f4d345985e65a1757a07f2/fastai/torch_core.py#L220


            else: l1 += list(trainable_params(c))
        #Since we scan the children separately, we might get duplicates (tied weights). We need to preserve the order
        #for the optimizer load of state_dict
        l1,l2 = uniqueify(l1),uniqueify(l2)
        split_params += [l1, l2]
    return split_params


def set_bn_eval(m:nn.Module)->None:
    "Set bn layers in eval mode for all recursive children of `m`."
    for l in m.children():
        if isinstance(l, bn_types) and not next(l.parameters()).requires_grad:
            l.eval()
        set_bn_eval(l)


def batch_to_half(b:Collection[Tensor])->Collection[Tensor]:
    "Set the input of batch `b` to half precision."
    return [to_half(b[0]), b[1]]


def bn2float(module:nn.Module)->nn.Module:
    "If `module` is batchnorm don't use half precision."
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm): module.float()