Siddhesh Sali

← All work

ml

Transfer Learning on Object Detection Models

Retraining SSD-MobileNet detectors on small custom datasets, and measuring how much of the feature extractor is worth freezing — the answer was the first four blocks.

A diagram: nine vertical bands in a row, the first five hatched and dark, the rest solid lilac and brightening to the right, over two amber bars where the lower one is less than half the length of the upper.

An internship at Shunya IoT, summer 2021, on the practical question behind most applied computer vision: you have a pretrained model, a few hundred labelled images of something it has never seen, and a Colab session that will time out. What do you actually do?

The warm-up was a binary classifier — MobileNetV2 pretrained on ImageNet, fine-tuned from layer 120 at a learning rate of 1e-4, reaching about 97% validation accuracy in eight epochs. That part is a recipe, and it works.

Detectors are less forgiving

Detection is harder than classification, and TensorFlow’s Object Detection API makes it harder still by hiding the training loop. Two detectors came out of it: a personal-protective-equipment detector on ssd_mobilenet_v2_fpnlite_640x640 — mask, face shield, gloves, goggles, full cover — and a firearms detector on ssd_mobilenet_v1_fpn_640x640, from a public Kaggle set. Both reached a total loss below 0.25 after 4,000 steps, about two hours of training.

Two hours per experiment is the constraint that made the rest of this project interesting. When each run costs an afternoon, the question stops being “which model is best” and becomes “what can I stop computing”.

How much to freeze

The obvious candidate is the feature extractor. Its early blocks learned edges and colour on ImageNet, and edges do not change when the subject does — so their gradients are arguably work you are paying for and throwing away.

So I ran it: the same model, on the same dataset, with the first n MobileNet-V1 blocks frozen, for n in nothing, four, eight, and all of them, and timed the steps.

Frozen Seconds per step
None (first run) 1.88
None (second run) 1.805
First 4 blocks 0.828
First 8 blocks 0.844
Whole extractor 0.845

The interesting column is the one that isn’t there. Freezing four blocks more than doubles the speed; freezing four more on top of that buys nothing — 0.844 against 0.828, which is noise. Whatever is expensive about the backward pass is concentrated in the first few blocks, and after that you are freezing layers that were cheap anyway while giving up their ability to adapt.

Plotted against wall-clock time rather than against steps, the frozen runs reach a lower loss and reach it sooner. The best run of the five froze four blocks; the worst froze none.

Freeze the first four to eight blocks. On this model and this size of dataset, that is the whole finding — and it is the kind of finding that only shows up if you time the runs instead of counting the steps.

Gallery

0404
A plot of training loss against wall-clock time for four runs. The two frozen runs, in blue, reach a lower loss well before the unfrozen runs in red and pink.
Fig. 01Loss against wall-clock time, not against steps. The frozen runs finish lower and sooner.
A plot comparing the loss curves for freezing four blocks and freezing eight. The two curves lie almost on top of each other.
Fig. 02Four frozen against eight. Almost the same curve — so take the one that costs less per step.
Training and validation accuracy rising to about 0.97 over eight epochs, with the matching cross-entropy loss falling below 0.15.
Fig. 03The classifier warm-up: MobileNetV2 fine-tuned from layer 120, eight epochs.
A three-by-three grid of validation images — camels, armadillos, ostriches, llamas and alpacas — each labelled with the model’s prediction and the true label.
Fig. 04The confusions that remain are the ones you would expect: a llama is very nearly an alpaca.