Fine-tune a pretrained fastai model without destroying it: freeze/unfreeze schedule and discriminative learning rates

TL;DR

When full unfreezing (A) keeps destroying the body even at low lr, common with small datasets or when the pretrained domain is far from yours, unfreeze progressively, ULMFiT-style: `python learn.freeze() learn.fit\_one\_cycle(3, 1e-2) # head only learn.freeze\_to(-2) learn.fit\_one\_cycle(3, slice(1e-2/2.64, 1e-2)) # last two groups learn.freeze\_to(-3) learn.fit\_one\_cycle(3, slice(5e-3/2.64, 5e-3)) # last three groups learn.unfreeze() learn.fit\_one\_cycle(4, slice(1e-3/2.6**4, 1e-3)) # everything Each freeze_to` clears optimizer state (fact 1), so treat every phase as a fresh optimizati

When to use

You are seeing this: Discriminative learning rates are a slice, spread log-evenly. fitonecycle(5, slice(1e-6, 1e-3)) gives the earliest layer group 1e-6, the last group 1e-3, and the middle groups log-spaced values in between (even_mults over the parameter groups).... Use this skill when you run into "Fine-tune a pretrained fastai model without destroying it: freeze/unfreeze schedule and discriminative learning rates".

When not to use

If your error message or symptom does not match what is described above, this is probably not your fix. Search for your exact error text instead of forcing this one to fit.

Versions

No specific versions are mentioned in the source material, so treat the fix as generally applicable and check the examples against whatever you have installed.

Use this when fine-tuning a pretrained model (vision_learner, text_classifier_learner, or any Learner with a pretrained backbone) goes wrong after unfreezing: validation loss spikes, the model overfits immediately, or unfreezing buys nothing over the frozen head. Classify the failure from the loss curves first, then apply the matching schedule fix.

0. The mechanics you need

learn.freeze()      # freeze_to(-1): train only the last layer group (the head)
learn.unfreeze()    # freeze_to(0):  train every layer group
learn.freeze_to(-2) # freeze all but the last two layer groups (gradual unfreezing)

Three facts that cause most of the confusion:

1. freeze/unfreeze/freeze_to clear the optimizer state. freeze_to calls opt.clear_state(), momentum and second-moment estimates are wiped every time you change what's frozen. Don't try to carry optimizer momentum across an unfreeze boundary; plan the schedule as separate phases instead.

2. Discriminative learning rates are a slice, spread log-evenly. fit_one_cycle(5, slice(1e-6, 1e-3)) gives the earliest layer group 1e-6, the last group 1e-3, and the middle groups log-spaced values in between (even_mults over the parameter groups). slice(1e-3) alone (no start) means [1e-4, ..., 1e-4, 1e-3], every group except the last gets stop/10. Early groups (generic features) always want the small end; the head wants the large end.

3. fit_one_cycle's lr is the peak, not the starting rate. fit_one_cycle(n_epoch, lr_max, div=25., div_final=1e5, pct_start=0.25) warms up from lr_max/25 to lr_max, then cosine-anneals down to lr_max/1e5. So a suggested lr of 3e-3 peaks at 3e-3 mid-training.

The standard recipe (frozen head first, then discriminative unfreeze):

learn.freeze()
lr = learn.lr_find().valley          # suggestion named tuple; .valley/.minimum/.steep/.slide
learn.fit_one_cycle(4, lr)

learn.unfreeze()
lr2 = learn.lr_find().valley         # re-run: different layer groups, different curve
learn.fit_one_cycle(6, slice(lr2/100, lr2))

Re-running lr_find() after unfreezing is not optional hygiene, the unfrozen model has different layer groups and a different loss landscape, so the frozen-phase lr is the wrong answer for the body.

1. Classify the failure from the curves

Plot with learn.recorder.plot_loss() and compare against the frozen phase. Note: in fastai v2 the monitored value is valid_loss (not the v1 name val_loss), old forum snippets using val_loss will silently monitor nothing.

A. valid_loss spikes immediately after unfreezing and never recovers

The learning rate is too high for the early layer groups, the pretrained body is being destroyed, not fine-tuned. Observable: the very first unfrozen epoch is much worse than the best frozen epoch.

• Recover: reload the best frozen weights. If you trained the frozen phase with SaveModelCallback(monitor='valid_loss', fname='head'), learn.load('head') gets them back; otherwise retrain the head.

• Redo the unfreeze with a discriminative slice whose low end is 10 100x below the frozen lr: learn.fit_one_cycle(4, slice(lr/100, lr)). Passing a single float after unfreeze() applies the same lr to every group, including the early ones, which is exactly what destroys them.

B. Frozen-phase valid_loss was still decreasing when you unfroze

You unfroze early. An undertrained head plus an unfrozen body mostly buys overfitting, not accuracy (a recurring finding on the fastai forums: once the frozen model stops improving, try unfreezing; while it's still improving, don't).

• Fix: go back to learn.freeze() and train the head until valid_loss plateaus, then unfreeze. Judge "plateau" from the recorder values, not from a fixed epoch count.

C. Unfrozen valid_loss improves for an epoch or two, then rises while train loss keeps falling

Classic overfitting, the unfrozen model has far more capacity than the frozen head. Observable: train/valid curves diverge.

• Train fewer unfrozen epochs; the unfrozen phase usually needs fewer epochs than the frozen phase.

• Add EarlyStoppingCallback(monitor='valid_loss', patience=2) and SaveModelCallback(monitor='valid_loss') so the best weights survive even if you overshoot:

from fastai.callback.tracker import SaveModelCallback, EarlyStoppingCallback
learn.fit_one_cycle(8, slice(lr/100, lr),
                    cbs=[SaveModelCallback(monitor='valid_loss'),
                         EarlyStoppingCallback(monitor='valid_loss', patience=2)])

• Increase weight decay (wd argument to fit_one_cycle) before reaching for a smaller model.

D. Gradual unfreezing: unfreezing one group at a time

When full unfreezing (A) keeps destroying the body even at low lr, common with small datasets or when the pretrained domain is far from yours unfreeze progressively, ULMFiT-style:

learn.freeze()
learn.fit_one_cycle(3, 1e-2)                 # head only
learn.freeze_to(-2)
learn.fit_one_cycle(3, slice(1e-2/2.6**4, 1e-2))  # last two groups
learn.freeze_to(-3)
learn.fit_one_cycle(3, slice(5e-3/2.6**4, 5e-3)) # last three groups
learn.unfreeze()
learn.fit_one_cycle(4, slice(1e-3/2.6**4, 1e-3)) # everything

Each freeze_to clears optimizer state (fact 1), so treat every phase as a fresh optimization run with its own lr. Lower the slice's top end as more groups join, the earlier groups must move less.

2. Pre-flight checks before any unfreeze

1. len(learn.opt.param_lists), know how many layer groups the slice will spread across. A vision_learner resnet34 has 3; a single float lr after unfreeze hits all of them equally.

2. Frozen valid_loss has plateaued (failure B).

3. SaveModelCallback(monitor='valid_loss') is attached so a bad unfreeze is recoverable with learn.load('model'), the default filename is 'model' in learn.path/learn.model_dir.

4. The unfrozen lr comes from an lr_find() run after unfreeze(), not from the frozen phase.

3. Checklist

1. Always train the head frozen first; unfreeze only on a plateau.

2. After unfreeze(), re-run lr_find() and use a slice(lo, hi) discriminative lr, never a single float for all groups.

3. Optimizer state resets on every freeze/unfreeze; schedule in independent phases.

4. Guard every unfrozen run with SaveModelCallback + early stopping on valid_loss (v2 name, not val_loss).

5. If full unfreezing destroys the body, go gradual with freeze_to(-n).

Built from the fastai v2 source: Learner.freeze/freezeto/unfreeze and opt.clearstate() in fastai/learner.py, Optimizer.sethyper slice handling (evenmults log spacing) in fastai/optimizer.py, fitonecycle (div=25., divfinal=1e5, pctstart=0.25) in fastai/callback/schedule.py, SaveModelCallback/EarlyStoppingCallback (monitor='valid_loss') in fastai/callback/tracker.py, and recurring forums.fast.ai threads on frozen-vs-unfrozen training and overfitting after unfreeze.