Layer Freezing

The problem, from first principles

The cheapest form of transfer learning is to freeze most of a pretrained network and train only the top. Frozen parameters receive no gradient updates, which saves memory (no optimizer state), time (no backward pass through frozen early layers is needed for their weights) and reduces overfitting when the new dataset is small. The choice is a spectrum: freeze everything except a new head (feature extraction), freeze only the first few blocks, or freeze nothing. Deciding which parameters are trainable from their names is a few lines of logic, and the number of trainable parameters is the headline number reported for any fine-tuning recipe.

From theory to code

Implement trainable_mask and count_trainable.

Constraints

• Parameters are named embed.<...>, layers.<i>.<...> or head.<...>. trainable_mask(names, n_frozen) returns a list of booleans: embed and layers.i for i < n_frozen are frozen (False), everything else (layers.i with i >= n_frozen, and head) is trainable (True). Any other name is trainable.
• count_trainable(sizes, mask) is the sum of sizes[k] over the entries where mask[k] is true, as an int, and fraction_trainable(sizes, mask) is that count divided by the total number of parameters.

Hints