KnowledgeDistillation Callback

Train a student model with a teacher using KnowledgeDistillationCallback 
from fastai.vision.all import *
from fasterai.core.all import *
from fasterai.distill.all import *

Distill a teacher into a student

KnowledgeDistillationCallback blends a distillation term into the student’s loss: at every batch it runs the teacher on the same inputs, compares the two models with the loss function you pass, and interpolates between that term and the ground-truth loss with weight (so at weight=0.9 the ground-truth term carries 0.1). This page trains one teacher (ResNet-34) and three students (ResNet-18), all at 64 px on the PETS cat/dog task.

1. Data

path = untar_data(URLs.PETS)
files = get_image_files(path/"images")

def label_func(f): return f[0].isupper()

dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(64))

2. The teacher

A pretrained ResNet-34, fine-tuned for ten epochs:

teacher = vision_learner(dls, resnet34, metrics=accuracy)
teacher.unfreeze()
teacher.fit_one_cycle(10, 1e-3)
epoch train_loss valid_loss accuracy time
0 0.746749 0.465374 0.878890 00:04
1 0.456684 0.337743 0.851150 00:03
2 0.306100 0.256743 0.896482 00:03
3 0.230790 0.197438 0.924222 00:03
4 0.175541 0.999262 0.843031 00:03
5 0.162185 7.491956 0.633965 00:03
6 0.114214 0.193958 0.928281 00:04
7 0.075693 0.171177 0.936401 00:04
8 0.038465 0.169076 0.941137 00:04
9 0.022227 0.171598 0.943843 00:04 
from math import sqrt

def report(learn, name):
    "Validation accuracy with its Wilson 95% interval"
    n = len(learn.dls.valid_ds)
    with learn.no_bar(): acc = float(learn.validate()[1])
    k, z = round(acc*n), 1.96
    p, d = k/n, 1 + z**2/n
    c = p + z**2/(2*n)
    h = z*sqrt(p*(1-p)/n + z**2/(4*n**2))
    print(f"{name}: {acc:.2%} ({k}/{n}), Wilson 95% [{(c-h)/d:.2%}, {(c+h)/d:.2%}]")

report(teacher, "teacher, ResNet-34")
teacher, ResNet-34: 94.38% (1395/1478), Wilson 95% [93.09%, 95.45%]

3. Student, no distillation

Same data, epochs and learning rate as the teacher; the students carry no pretrained weights:

student = Learner(dls, resnet18(num_classes=2), metrics=accuracy)
student.fit_one_cycle(10, 1e-3)
epoch train_loss valid_loss accuracy time
0 0.619088 0.605188 0.689445 00:03
1 0.586678 0.566887 0.706360 00:03
2 0.546209 0.610709 0.667794 00:03
3 0.513684 0.498844 0.747632 00:03
4 0.463067 0.586233 0.752368 00:03
5 0.403649 0.653835 0.764547 00:03
6 0.356432 0.457676 0.814614 00:02
7 0.277365 0.373427 0.837618 00:02
8 0.213076 0.394300 0.828146 00:02
9 0.160289 0.389774 0.838972 00:02 
report(student, "student alone")
student alone: 83.90% (1240/1478), Wilson 95% [81.94%, 85.68%]

4. SoftTarget

SoftTarget compares the softened predictions of the two models. schedule=cos makes the weight of that term follow a cosine progression over training:

student = Learner(dls, resnet18(num_classes=2), metrics=accuracy)
kd = KnowledgeDistillationCallback(teacher.model, SoftTarget, schedule=cos)
student.fit_one_cycle(10, 1e-3, cbs=kd)
epoch train_loss valid_loss accuracy time
0 0.627812 0.699146 0.635995 00:02
1 0.697857 0.745051 0.706360 00:03
2 0.847422 0.836012 0.741543 00:02
3 0.993758 0.988241 0.744926 00:02
4 1.140755 1.061720 0.788227 00:02
5 1.200129 1.146610 0.804465 00:02
6 1.133657 0.959932 0.847091 00:03
7 0.952126 0.933980 0.855210 00:02
8 0.804851 0.941943 0.859269 00:02
9 0.694702 0.928969 0.866712 00:03 
report(student, "student + SoftTarget")
student + SoftTarget: 86.67% (1281/1478), Wilson 95% [84.84%, 88.31%]

5. Attention, with the layers matched automatically

Intermediate-feature losses need pairs of layers, one in each model. match_feature_layers builds them: it runs one forward pass through each model, groups layers by output (H, W), and picks a match at each resolution.

student_model = resnet18(num_classes=2)

pairs = match_feature_layers(student_model, teacher.model, torch.randn(1, 3, 64, 64))
print(pairs)

student = Learner(dls, student_model, metrics=accuracy)
kd = KnowledgeDistillationCallback(
    teacher.model, Attention,
    pairs['student'], pairs['teacher'],
    weight=0.9
)
student.fit_one_cycle(10, 1e-3, cbs=kd)
{'student': ['conv1', 'layer1', 'layer2', 'layer3', 'layer4'], 'teacher': ['0.0', '0.4', '0.5', '0.6', '0']}
epoch train_loss valid_loss accuracy time
0 0.087690 0.083192 0.700271 00:03
1 0.079434 0.069559 0.728011 00:02
2 0.069988 0.071774 0.747632 00:03
3 0.061195 0.056563 0.803112 00:04
4 0.054241 0.053276 0.811231 00:04
5 0.045464 0.052328 0.815968 00:04
6 0.038911 0.043625 0.870095 00:04
7 0.030366 0.041490 0.872124 00:04
8 0.022930 0.041212 0.884980 00:05
9 0.019471 0.040854 0.884980 00:06 
report(student, "student + Attention")
student + Attention: 88.50% (1308/1478), Wilson 95% [86.77%, 90.03%]

The teacher is a fastai vision_learner model, so its layers are named positionally: '0.0' is the body’s first convolution, '0.4', '0.5', '0.6' are layer1, layer2, layer3. The last pair maps the student’s layer4 onto '0', the body itself — the body’s output is the output of its last block, so this hooks the same tensor as '0.7' would.

Scope. PETS cat/dog labels, images resized to 64 px, ten epochs per arm at lr=1e-3; single run, no seed fixed. Every accuracy above is measured on the same 1478 validation images and printed with its Wilson 95% interval. The teacher is an ImageNet-pretrained ResNet-34 with a fastai head, the three students are resnet18(num_classes=2) trained from scratch — pretraining, not architecture, is the dominant difference between the first number and the other three. Each distillation arm also changes more than one thing with respect to the student trained alone: SoftTarget sets schedule=cos, Attention sets weight=0.9. So the four numbers show that the callback runs and what it produced here, not the size of an effect. This page measures no latency.

Summary

Tool / Function What it gives you
KnowledgeDistillationCallback(teacher, loss, activations_student, activations_teacher, weight, schedule) A distillation term added to the student’s loss during fit
match_feature_layers(student, teacher, x) Layer pairs matched by output resolution, for the feature-based losses
SoftTarget Compares the two models’ softened predictions
Attention Compares attention maps taken at the matched layers

The losses exported by fasterai.distill.losses are SoftTarget, Logits, Mutual, DecoupledKD, Attention, ActivationBoundaries, FitNet and Similarity.

See Also